diff --git a/.github/workflows/release.yml b/.github/workflows/release.yml index 5032a1ed..685feac8 100644 --- a/.github/workflows/release.yml +++ b/.github/workflows/release.yml @@ -203,6 +203,16 @@ jobs: echo "publish_version=${PUBLISH_VERSION}" >> $GITHUB_OUTPUT echo "Tag ${GITHUB_REF_NAME} @ ${SHORT_SHA_HEX} -> TestPyPI version ${PUBLISH_VERSION}" + # The packaged weightslab/VERSION states the version being published, so + # the artifact agrees with itself before it is uploaded (the commit back + # to dev/main is sync-version-file's job, after the release). + - name: Write weightslab/VERSION + env: + PUBLISH_VERSION: ${{ steps.compute_version.outputs.publish_version }} + run: | + printf '%s\n' "${PUBLISH_VERSION}" > weightslab/VERSION + echo "weightslab/VERSION -> $(cat weightslab/VERSION)" + - name: Build distributions env: SETUPTOOLS_SCM_PRETEND_VERSION_FOR_WEIGHTSLAB: ${{ steps.compute_version.outputs.publish_version }} @@ -608,7 +618,24 @@ jobs: - name: Install build tools run: python -m pip install --upgrade pip build twine + # The packaged weightslab/VERSION states the version being published: + # the tag's own, PEP 440-normalized (v2.0.3 -> 2.0.3). The commit back to + # main/dev is sync-version-file's job, after the release. + # + # The build is then pinned to that same version. Writing the file dirties + # the checkout, and setuptools_scm's guess-next-dev turns a dirty tree at + # a tag into the NEXT dev version (2.0.4.dev0) instead of the tag's. + - name: Write weightslab/VERSION + id: release_version + run: | + VERSION="$(python -c 'import os; from packaging.version import Version; print(Version(os.environ["GITHUB_REF_NAME"]))')" + printf '%s\n' "${VERSION}" > weightslab/VERSION + echo "version=${VERSION}" >> $GITHUB_OUTPUT + echo "weightslab/VERSION -> ${VERSION}" + - name: Build distributions + env: + SETUPTOOLS_SCM_PRETEND_VERSION_FOR_WEIGHTSLAB: ${{ steps.release_version.outputs.version }} run: python -m build - name: Capture version from wheel @@ -931,3 +958,95 @@ jobs: files: dist/* env: GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }} + + # ═══════════════════════════════════════════════════════════════════════════ + # VERSION FILE SYNC — after either release, commit the tag's version into + # weightslab/VERSION: on main AND dev for a main release, on dev only for a + # dev release (a dev version never lands on main this way). + # ═══════════════════════════════════════════════════════════════════════════ + # weightslab/VERSION is the last-resort source of weightslab.__version__ (and + # so of the telemetry version): read only when there is no git metadata, no + # setuptools_scm-built _version.py and no installed dist metadata. It never + # overrides a git-derived version, so builds and checkouts are unaffected. + # + # main and dev both reject direct pushes (PR + status-check rulesets) unless + # the pusher is a ruleset bypass actor, which GITHUB_TOKEN is not — hence + # VERSION_SYNC_TOKEN: a token of a bypass actor with contents:write on this + # repo. "[skip ci]" stops that push from re-triggering ci.yml/docs-pages.yml + # (a PAT push triggers workflows, a GITHUB_TOKEN push would not), which would + # otherwise run a TestPyPI publish for the push to main. + sync-version-file: + name: Sync weightslab/VERSION + needs: [detect-target, create-dev-release, create-main-release] + # !cancelled(): exactly one of the two release jobs is skipped on any tag. + if: ${{ !cancelled() && (needs.create-main-release.result == 'success' || needs.create-dev-release.result == 'success') }} + runs-on: ubuntu-latest + permissions: + contents: read + steps: + - name: Require the push token + env: + VERSION_SYNC_TOKEN: ${{ secrets.VERSION_SYNC_TOKEN }} + run: | + if [ -z "${VERSION_SYNC_TOKEN}" ]; then + echo "::error::VERSION_SYNC_TOKEN is not set; main/dev reject pushes from GITHUB_TOKEN, so weightslab/VERSION was not updated for ${GITHUB_REF_NAME}. The release itself is published." + exit 1 + fi + + - uses: actions/setup-python@v6 + with: + python-version: '3.11' + + # PEP 440-normalized, so the file holds what a build of the tag reports: + # v2.0.2 -> 2.0.2, v1.3.1-dev0 -> 1.3.1.dev0. + - name: Compute version and target branches + id: target + env: + IS_MAIN: ${{ needs.detect-target.outputs.is_main }} + run: | + set -euo pipefail + python -m pip install --quiet packaging + VERSION="$(python -c 'import os; from packaging.version import Version; print(Version(os.environ["GITHUB_REF_NAME"]))')" + BRANCHES="dev" + if [ "${IS_MAIN}" = "true" ]; then + BRANCHES="main dev" + fi + echo "version=${VERSION}" >> $GITHUB_OUTPUT + echo "branches=${BRANCHES}" >> $GITHUB_OUTPUT + echo "Tag ${GITHUB_REF_NAME} -> weightslab/VERSION=${VERSION} on: ${BRANCHES}" + + - uses: actions/checkout@v4 + with: + token: ${{ secrets.VERSION_SYNC_TOKEN }} + + # Each branch is rewritten from its fresh tip; a push rejected because the + # branch moved meanwhile is retried on the new tip rather than rebased. + - name: Commit weightslab/VERSION + env: + VERSION: ${{ steps.target.outputs.version }} + BRANCHES: ${{ steps.target.outputs.branches }} + run: | + set -euo pipefail + git config user.name "github-actions[bot]" + git config user.email "41898282+github-actions[bot]@users.noreply.github.com" + for BRANCH in ${BRANCHES}; do + for attempt in 1 2 3; do + git fetch --depth=1 origin "${BRANCH}" + git checkout -q -B "${BRANCH}" FETCH_HEAD + printf '%s\n' "${VERSION}" > weightslab/VERSION + git add weightslab/VERSION + if git diff --cached --quiet; then + echo "${BRANCH}: weightslab/VERSION is already ${VERSION}" + continue 2 + fi + git commit -q -m "release: weightslab/VERSION -> ${VERSION} [skip ci]" + if git push origin "HEAD:${BRANCH}"; then + echo "${BRANCH}: weightslab/VERSION -> ${VERSION}" + continue 2 + fi + echo "Push to ${BRANCH} rejected (attempt ${attempt}/3); retrying on its new tip" + sleep 5 + done + echo "::error::Could not push weightslab/VERSION to ${BRANCH}" + exit 1 + done diff --git a/AGENTS.md b/AGENTS.md index a6483ff2..254b42c7 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -129,6 +129,9 @@ Conventions that matter for correctness: and pause/resume or stats will misbehave.** - Use `model.get_age()` (steps actually trained; survives checkpoint reloads), not the raw loop counter. +- Pause from code with `wl.pause_training()` (the counterpart of + `wl.start_training()`): the next guarded step waits for Play. Running a cell + in the Studio notebook pauses training the same way, and it stays paused. - `task_type` on the dataset/model selects rendering: `classification`, `segmentation`, `detection`, `detection_pointcloud`. - **Hyperparameter handle access:** the registered hyperparameters proxy @@ -143,6 +146,7 @@ Recording values (pick by what the value is *about*, not by convenience): | one annotation | `wl.save_instance_signals(...)` | one box's IoU | | a group of samples | `wl.save_group_signals(signals={...}, group_ids=[...])` | a pair's contrastive loss | | one training **step** | `wl.save_model_signals(signals={...})` | a gradient norm | +| a sample's place in a 2-D/3-D **projection** | `wl.save_projection_coords(coords, batch_ids=ids, prefix="tsne")` | your own t-SNE / umap-learn layout | The first three write dataframe rows (sortable/filterable in the grid); the fourth only plots a curve. Never fake a step-level value by broadcasting it @@ -189,6 +193,7 @@ ones when debugging: | `GRPC_MAX_MESSAGE_BYTES` | `268435456` (256 MB) | Raise it if large tensors/image batches fail. | | `WEIGHTSLAB_DISABLE_WATCHDOGS` | `0` | Set `1` when debugging with breakpoints (see §5). | | `GRPC_WATCHDOG_STUCK_SECONDS` | `60` | Lock/RPC stuck threshold + lock-acquire timeout. | +| `WEIGHTSLAB_PROJECTION` | on | Live 3-D parametric-UMAP projection (beta). `0` removes it; per model: `watch_or_edit(model, flag="model", projection=False)`. Tuning: `_EVERY`, `_DIM`, `_NEIGHBORS`, `_GRAPH` — see `docs/projection.rst`. | **Frontend (Weights Studio) — runtime-injected `window.*` globals:** @@ -265,6 +270,14 @@ The train step isn't wrapped in `guard_training_context` (or eval in `guard_testing_context`). See §3 — these context managers are how the system gates and separates phases. +**Projection board stays empty, or stopped updating.** `wl.projection.get_tracker()` +is `None` when it is off (`WEIGHTSLAB_PROJECTION=0` or `projection=False`). +Otherwise `.stats()` says why: `fits == 0` means the hooked layer is not on the +forward path (pass `projection={"layer": ""}`); +`disabled_reason` is set when 5 fits in a row failed and it turned itself off; +`skipped_nonfinite > 0` means the model's features went NaN. Training is never +affected by any of these. Full page: `weightslab/docs/projection.rst`. + **Large weights/images fail to transfer.** Raise `GRPC_MAX_MESSAGE_BYTES`. **The agent bar says it's unconfigured.** The LLM agent needs a provider: a diff --git a/CHANGELOG.md b/CHANGELOG.md index 384e44f5..c53073e6 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1 +1 @@ -# Changelog - 2026-10-01 v2.0.2 +# Changelog - 2026-10-09 v2.0.3.dev2 diff --git a/docs/_static/custom.css b/docs/_static/custom.css index 48647898..4f83055b 100644 --- a/docs/_static/custom.css +++ b/docs/_static/custom.css @@ -110,24 +110,6 @@ body[data-theme="dark"] .wl-only-dark { background: rgba(120, 120, 120, 0.12); } -/* Sandbox CTA -- warm gradient built on the #f0a500 accent used elsewhere in - the docs chrome, so the external "try it now" link reads as distinct from - the in-docs Quickstart/Examples buttons without leaving the palette. */ -.wl-hero-cta-sandbox, -.wl-hero-cta-sandbox:link, -.wl-hero-cta-sandbox:visited, -.wl-hero-cta-sandbox:active, -.wl-hero-cta-sandbox:focus { - background: linear-gradient(135deg, #f0a500 0%, #e2681c 100%) !important; - border-color: #e2871a !important; - color: #ffffff !important; -} - -.wl-hero-cta-sandbox:hover { - background: linear-gradient(135deg, #ffb61f 0%, #f07a2b 100%) !important; - border-color: #f0951f !important; -} - .wl-version-switcher { display: inline-flex; align-items: center; diff --git a/docs/_static/examples-gallery.js b/docs/_static/examples-gallery.js index a09ab4ac..1730247d 100644 --- a/docs/_static/examples-gallery.js +++ b/docs/_static/examples-gallery.js @@ -22,6 +22,13 @@ url: 'examples/pytorch/classification.html', colab: COLAB + 'PyTorch/wl-classification.ipynb' }, + { + badge: 'PyTorch', color: 'pytorch', + title: 'Classification, Fashion-MNIST + UMAP', + desc: 'Train with the live parametric UMAP, reload the 1000-step weights, discard the lowest-loss samples, then compare a PCA you write yourself with the UMAP.', + tags: ['classification', 'fashion-mnist', 'projection', 'umap', 'pca', 'checkpoint', 'beta'], + url: 'examples/pytorch/fashion_mnist_umap.html' + }, { badge: 'PyTorch', color: 'pytorch', title: 'Segmentation, BDD100k', @@ -90,6 +97,13 @@ desc: 'Per-step training dynamics: global and per-layer gradient norms, weight norms and activation statistics, from one argument on the model wrap.', tags: ['model signals', 'gradient norm', 'activations', 'per-layer', 'training dynamics'], url: 'examples/usecases/model_signals.html' + }, + { + badge: 'Usecase', color: 'usecase', + title: 'Projections, Fashion-MNIST', + desc: 'The built-in live UMAP while training, then a run that pauses itself at convergence and plugs a notebook-written t-SNE into the Projection Board.', + tags: ['projection', 'umap', 't-sne', 'notebook', 'beta'], + url: 'examples/usecases/projection_fashion_mnist.html' } ]; diff --git a/docs/_static/screenshots/projection-board.png b/docs/_static/screenshots/projection-board.png new file mode 100644 index 00000000..31bffe85 Binary files /dev/null and b/docs/_static/screenshots/projection-board.png differ diff --git a/docs/agent.rst b/docs/agent.rst index d4caf4fb..f9385754 100644 --- a/docs/agent.rst +++ b/docs/agent.rst @@ -166,7 +166,7 @@ Remote provider, OpenRouter - Ask for a schema-validated plan instead of free-form JSON + repair. More reliable, but only on routes that support JSON-schema output (Gemini, GPT-4o). * - ``fallback_to_local`` - - — + - (none) - ``true`` (the shipped ``agent_config.yaml`` sets ``false``) - Also set up Ollama, so a failing/absent cloud key still leaves a working agent. diff --git a/docs/build_docs.sh b/docs/build_docs.sh index 3ea30acd..18e626ab 100644 --- a/docs/build_docs.sh +++ b/docs/build_docs.sh @@ -4,17 +4,42 @@ set -euo pipefail ROOT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)" cd "$ROOT_DIR" +# The interpreter: PYTHON_BIN, else the activated venv, else the first Python on +# PATH that can build -- one that already has Sphinx, or has pip to install it. +# +# The .exe names matter on Windows: from PowerShell, `bash` is WSL's bash, where +# python3 is often a Linux Python with no pip, while the venv activated in that +# PowerShell is reachable only as python.exe (WSL runs Windows programs by their +# full name). Every path handed to the interpreter below is relative to +# ROOT_DIR, so a Windows Python never sees a /mnt/c/... path. +can_build() { + "$@" -c "import sphinx, furo, myst_parser, sphinx_design, sphinxcontrib.mermaid" >/dev/null 2>&1 || + "$@" -m pip --version >/dev/null 2>&1 +} + +PYTHON_CMD=() if [[ -n "${PYTHON_BIN:-}" ]]; then PYTHON_CMD=("$PYTHON_BIN") -elif command -v python3 >/dev/null 2>&1; then - PYTHON_CMD=(python3) -elif command -v python >/dev/null 2>&1; then - PYTHON_CMD=(python) -elif command -v py >/dev/null 2>&1; then - PYTHON_CMD=(py -3) else - echo "[weightslab-docs] ERROR: no Python interpreter found in PATH." - echo "[weightslab-docs] Set PYTHON_BIN explicitly and retry." + candidates=() + if [[ -n "${VIRTUAL_ENV:-}" ]]; then + candidates+=("$VIRTUAL_ENV/bin/python" "$VIRTUAL_ENV/Scripts/python.exe") + fi + candidates+=("python3" "python" "py -3" "python.exe" "py.exe -3") + for candidate in "${candidates[@]}"; do + read -r -a cmd <<< "$candidate" + if command -v "${cmd[0]}" >/dev/null 2>&1 && can_build "${cmd[@]}"; then + PYTHON_CMD=("${cmd[@]}") + break + fi + done +fi + +if [[ ${#PYTHON_CMD[@]} -eq 0 ]]; then + echo "[weightslab-docs] ERROR: no Python with Sphinx or pip found in PATH" + echo "[weightslab-docs] (tried python3, python, py, python.exe, py.exe)." + echo "[weightslab-docs] Set PYTHON_BIN to one and retry, e.g.:" + echo " PYTHON_BIN=/path/to/venv/bin/python bash docs/build_docs.sh" exit 1 fi @@ -69,8 +94,22 @@ fi DOCS_HOST="${WEIGHTSLAB_DOCS_HOST:-127.0.0.1}" DOCS_PORT_REQUESTED="${WEIGHTSLAB_DOCS_PORT:-8000}" +# Serving needs only the standard library. When the build ran on a Windows +# python.exe from WSL, a native Python serves instead if there is one: `kill` +# from this shell cannot reliably stop a Windows process, so Ctrl+C would leave +# the server running. +SERVE_CMD=("${PYTHON_CMD[@]}") +if [[ "${PYTHON_CMD[0]}" == *.exe ]]; then + for candidate in python3 python; do + if command -v "$candidate" >/dev/null 2>&1 && "$candidate" -c "import http.server" >/dev/null 2>&1; then + SERVE_CMD=("$candidate") + break + fi + done +fi + # Pick the requested port, or the first free one above it. -DOCS_PORT="$("${PYTHON_CMD[@]}" - "$DOCS_HOST" "$DOCS_PORT_REQUESTED" <<'PY' || true +DOCS_PORT="$("${SERVE_CMD[@]}" - "$DOCS_HOST" "$DOCS_PORT_REQUESTED" <<'PY' || true import socket import sys @@ -110,7 +149,7 @@ cleanup() { trap cleanup EXIT INT TERM echo "[weightslab-docs] Serving $HTML_DIR at $DOCS_URL" -"${PYTHON_CMD[@]}" -m http.server "$DOCS_PORT" --bind "$DOCS_HOST" --directory "$HTML_DIR" >/dev/null 2>&1 & +"${SERVE_CMD[@]}" -m http.server "$DOCS_PORT" --bind "$DOCS_HOST" --directory docs/_build/html >/dev/null 2>&1 & SERVER_PID=$! # Wait for the socket to accept connections before pointing a browser at it. @@ -119,7 +158,7 @@ for _ in $(seq 1 100); do if ! kill -0 "$SERVER_PID" 2>/dev/null; then break fi - if "${PYTHON_CMD[@]}" - "$DOCS_HOST" "$DOCS_PORT" <<'PY' >/dev/null 2>&1 + if "${SERVE_CMD[@]}" - "$DOCS_HOST" "$DOCS_PORT" <<'PY' >/dev/null 2>&1 import socket import sys diff --git a/docs/checkpointing.rst b/docs/checkpointing.rst index b0a28e9f..b3d01117 100644 --- a/docs/checkpointing.rst +++ b/docs/checkpointing.rst @@ -34,7 +34,7 @@ Directory structure ``root_log_dir`` comes from the top-level ``root_log_dir`` key in your hyperparameters config (or the ``root_log_dir=`` kwarg on ``wl.watch_or_edit(..., flag="model"/"data"/...)`` for a per-object override) -— see :doc:`configuration`. +, see :doc:`configuration`. The experiment hash -------------------- diff --git a/docs/configuration.rst b/docs/configuration.rst index 21ec14f5..830a2cf4 100644 --- a/docs/configuration.rst +++ b/docs/configuration.rst @@ -74,7 +74,7 @@ Accepted by every ``flag`` value. - ``None`` - Override the root directory for checkpoints and logs for this object only. Defaults to ``WEIGHTSLAB_ROOT_LOG_DIR``. May also point - at a parent directory that fans out into several experiment roots — + at a parent directory that fans out into several experiment roots, see :doc:`checkpointing`. * - ``skip_previous_auto_load`` - ``False`` @@ -440,13 +440,41 @@ Logging disable. Only the main process writes one, ``DataLoader`` workers and other spawned children keep a terminal-only logger. * - ``WEIGHTSLAB_ROOT_LOG_DIR`` - - *(temporary directory)* - - Experiment directory (checkpoints, reports, notebook, logs). Also what - ``weightslab start [DIR]`` exports. The session log starts in - ``/weightslab_logs/`` and, when it is unset, starts in a temporary - directory and is **moved** into ``/weightslab_logs/`` as - soon as the experiment's ``root_log_dir`` resolves, so the log always - ends up beside the checkpoints. The path is printed when the process exits. + - *(training script dir)* + - Root directory where training log snapshots are saved. + Defaults to a ``root_log_dir/`` folder next to your training script. + * - ``WEIGHTSLAB_PROJECTION`` + - *(on)* + - Live 3-D parametric-UMAP projection of the model's representation + (see :doc:`projection`, beta). On by default; set ``0``, ``false``, ``no`` + or ``off`` to remove it entirely , no hook, no encoder, no cost. Wins + over the per-model ``projection=`` keyword; to turn it off for one model + only, pass ``projection=False`` to ``wl.watch_or_edit(..., flag="model")``. + * - ``WEIGHTSLAB_PROJECTION_EVERY`` + - ``50`` + - Training steps between projection fits. Raise it to make the projection + cheaper, lower it to make it track a fast-moving representation. + * - ``WEIGHTSLAB_PROJECTION_DIM`` + - ``3`` + - Projection output dimensions. ``2`` is drawn on the ``z = 0`` plane. + * - ``WEIGHTSLAB_PROJECTION_NEIGHBORS`` + - ``15`` + - UMAP ``n_neighbors`` for the kNN graph. Clamped to the graph size minus one. + * - ``WEIGHTSLAB_PROJECTION_GRAPH`` + - ``512`` + - Samples the projection's kNN graph is built over: a rolling buffer of the + most recent training samples, so graph quality does not depend on your + batch size. Raise it if the cloud is one blob instead of clusters. + ``EVERY``, ``DIM``, ``NEIGHBORS`` and ``GRAPH`` can each be overridden per + model with ``projection={...}`` (see :doc:`projection`). + * - ``WEIGHTSLAB_PROJECTION_MAX_POINTS`` + - ``70000`` + - Points the Projection Board draws per view, at most 400,000. A bigger + projection is drawn as a subsample of this many points; zooming in spends + them on a smaller region. Read by the training process (it serves the + board), on every request: ``os.environ[...] = "150000"`` in the Studio + notebook applies at the next camera move. A browser can override it for + itself (see :doc:`weights_studio_ui/main_area`). * - ``AUDIT_LOG_FORMAT`` - ``json`` - Output format for audit logs tracking all user interactions through gRPC. @@ -863,7 +891,7 @@ server and the backend SDK agent share. These control where it lives. - Host the spawned agent server **binds** to. Loopback by default (the server has filesystem access and must not be reachable off-machine on a normal local run). Set to ``0.0.0.0`` when running in a container reached - over an SSH tunnel / published port, so the published port can reach it — + over an SSH tunnel / published port, so the published port can reach it, the URL handed to the browser stays ``127.0.0.1`` either way. See :ref:`studio-bridging`. * - ``WEIGHTSLAB_UI_TRUSTED_HOSTS`` @@ -1234,7 +1262,7 @@ agent. Each toggle **removes the area from the UI** (the elements are hidden) started), so a disabled area costs nothing at runtime. Like the bounding-box render limits, these are set as environment variables -before ``weightslab start`` and injected into ``config.js`` at startup — +before ``weightslab start`` and injected into ``config.js`` at startup, changing them needs no rebuild, just a restart + browser reload. For a local ``vite`` dev server, use the ``VITE_`` fallbacks shown below. Every toggle **defaults to enabled**; set it to ``0`` / ``false`` / ``no`` / ``off`` (any diff --git a/docs/custom_evaluation.rst b/docs/custom_evaluation.rst index a7606775..10085074 100644 --- a/docs/custom_evaluation.rst +++ b/docs/custom_evaluation.rst @@ -5,7 +5,7 @@ When the CLI's ``evaluate`` command (or the UI's evaluate action, or the agent) runs an evaluation, something has to actually iterate the loader and compute the numbers. By default WeightsLab builds that runner for you out of whatever signals you've already registered as ``flag="loss"``/``flag="metric"`` -— but decorating your own function with ``@wl.eval_fn`` replaces that default +, but decorating your own function with ``@wl.eval_fn`` replaces that default with your real evaluation logic. The built-in default diff --git a/docs/examples/index.rst b/docs/examples/index.rst index 2ec7c309..f8f56408 100644 --- a/docs/examples/index.rst +++ b/docs/examples/index.rst @@ -11,24 +11,6 @@ Browse all WeightsLab examples. Use the search bar or filter by framework.

Loading gallery…

-Try it without installing anything ------------------------------------ - -Prefer to look before you install? The sandbox is a hosted Weights Studio -running against a live experiment, with the same boards, plots, and agent -chat described throughout these docs. It opens in read-only demo mode, you -can browse, sort, filter, and inspect samples, but write actions (tagging, -discarding, training control, export) are disabled, so nothing you click can -break it. No account, no setup. - -.. raw:: html - -
- - Open Sandbox - -
- .. toctree:: :hidden: :maxdepth: 2 diff --git a/docs/examples/pytorch/fashion_mnist_umap.rst b/docs/examples/pytorch/fashion_mnist_umap.rst new file mode 100644 index 00000000..c28fe368 --- /dev/null +++ b/docs/examples/pytorch/fashion_mnist_umap.rst @@ -0,0 +1,169 @@ +Classification, Fashion-MNIST + UMAP (PyTorch) +=============================================== + +.. raw:: html + +
+ PyTorch + classification + fashion-mnist + projection + umap + pca + beta +
+ +**Example:** ``weightslab/examples/PyTorch/wl-fashion-mnist-umap/main.py`` + +**Task:** 10-class Fashion-MNIST classification with a small CNN, explored with +projections. The built-in parametric UMAP runs beside training; then you go back +to the weights of the first 1000 steps, discard the training samples that teach +the model least, lay what is left out with **a PCA you write yourself**, and +compare the two layouts in the :ref:`studio-projection-board` and in numbers. + +It is the plain-PyTorch counterpart of the +:doc:`projection use case <../usecases/projection_fashion_mnist>`, driven in +stages so you can look at the board between them. Background and every option of +the projection: :doc:`../../projection`. + +The stages +---------- + +1. Train with the live UMAP, keep a checkpoint +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +Nothing in the loop is projection code. Wrapping the model attaches the UMAP +(it is on by default); the dict only tunes it: + +.. code-block:: python + + model = wl.watch_or_edit( + FashionCNN().to(device), flag="model", device=device, + projection={"every_n_steps": 25, "graph_size": 1024}, # optional + ) + + train_to(1000) # the UMAP fits alongside + checkpoint_step = save_checkpoint() # weights + encoder + per-sample state + +The projection hooks the input of the last ``Linear`` (``fc2``): the 128-D +representation the classifier decides from, not the logits. The checkpoint holds +the weights, the UMAP encoder and every sample's loss at that step. + +2. Reload that checkpoint +~~~~~~~~~~~~~~~~~~~~~~~~~ + +.. code-block:: python + + manager = ledgers.get_checkpoint_manager() + manager.load_state(manager.get_current_experiment_hash(), target_step=checkpoint_step) + +Weights, optimizer, the UMAP encoder and each sample's loss and seen-count go +back to that step. Straight after stage 1 it changes nothing; it is the way back +to *these* weights if you train on and come back to try another cut. + +3. Discard what teaches little +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +.. code-block:: python + + table = sample_table() # one row per training sample + easy = lowest_loss_ids(table, 0.3) # the 30% with the lowest loss + wl.discard_samples(easy, discarded=True) # the loaders skip them from now on + +Samples with no loss yet are never picked: an unseen sample is not an easy one. +Expect the cut to come mostly from the clean classes (Trouser, Bag, the +footwear) and hardly at all from the tops the model still confuses. + +4. A PCA of your own +~~~~~~~~~~~~~~~~~~~~ + +.. code-block:: python + + def pca_coordinates(features, dim=2): # plain torch, an SVD + x = torch.as_tensor(features, dtype=torch.float32) + x = x - x.mean(0) + _, singular, axes = torch.linalg.svd(x, full_matrices=False) + return (x @ axes[:dim].T).numpy() + + wl.project_dataset(model, train_loader, method=pca_coordinates, + prefix="pca2d_step1000") + +``wl.project_dataset`` runs the loader through the model, collects the 128-D +features, hands them to **your function**, and stores what comes back as +``signals//pca2d_step1000_{x,y}``. The loader skips the samples you just +discarded, so the PCA covers what is left. The sweep always starts from the +loader's first batch (even though training was mid-epoch on it) and does not age +the model. Any algorithm returning ``(N, 2)`` or ``(N, 3)`` plugs in the same +way. The name carries the step, so snapshots stay apart in the picker. + +5. Compare +~~~~~~~~~~ + +In the Studio, *Data Exploration → Projection*, open the picker and switch +between ``umap`` (3-D, learned live) and ``pca2d_step1000`` (2-D, drawn on the +z = 0 plane); lasso the same region in both and compare what the grid loads. + +In the script, two more views of the same two layouts. A side-by-side picture is +written to ``/umap_vs_pca.png``. And ``knn_purity`` scores each +layout: for every point, the share of its *k* nearest neighbours **in the +picture** that have its class (1.0 = every neighbourhood is one class, about 0.1 += ten classes thrown together). The 128-D features are scored too, as the +reference both pictures summarize: + +.. code-block:: text + + neighbours sharing a point's class (k=10), over 14000 kept samples: + parametric UMAP (3-D) : 0.720 + your PCA (2-D) : 0.545 + the 128-D features : 0.858 (what both pictures summarize) + +*(One CPU run of the defaults; yours will differ in the last digits. UMAP keeps +three dimensions and PCA two, so read the gap as a rough size, not a verdict.)* + +6. Retrain on what is left +~~~~~~~~~~~~~~~~~~~~~~~~~~ + +The live UMAP **follows the model**: its encoder keeps fitting, so its layout +changes with the weights. The PCA does not: it is the snapshot you took, so to +see the PCA of the new model you project again under a new name +(``pca2d_step2000``). In the same run, 1000 more steps on the remaining samples +took the UMAP's score from 0.72 to 0.84 while a fresh PCA moved from 0.55 to +0.56: the representation got more class-separable, and only the layout that +learns shows it that clearly. + +Every projection you compute is stored with the checkpoint of that moment (see +:doc:`../../projection`), so restoring an older checkpoint brings back the +projections that existed then. + +Configuration +------------- + +``config.yaml`` carries the tuning of the UMAP (``projection:``) and of the +stages (``exploration:``): + +.. code-block:: yaml + + projection: + every_n_steps: 25 # fit the UMAP encoder every 25 training steps + graph_size: 1024 # samples in the kNN graph + exploration: + steps_initial: 1000 # train with the live UMAP, then keep a checkpoint + discard_fraction: 0.3 # share of training samples to drop, lowest loss first + steps_retrain: 1000 # train on what is left (0 = stop after the comparison) + knn_k: 10 # neighbours counted by the class-purity score + +Run it +------ + +.. code-block:: bash + + weightslab start # terminal 1: the studio, http://localhost:8080 + + cd weightslab/examples/PyTorch/wl-fashion-mnist-umap + python main.py --wait --keep-serving # terminal 2 + python main.py --steps 200 100 # a quick pass: 200 initial + 100 retrain steps + +``--wait`` stops for Enter after the comparison, so you can look at the board +before retraining; ``--keep-serving`` keeps the process (and the Studio +connection) up at the end. The defaults take a couple of minutes on a CPU. The +projection board is a **beta** feature. diff --git a/docs/examples/pytorch/index.rst b/docs/examples/pytorch/index.rst index e12e2e05..57a67b49 100644 --- a/docs/examples/pytorch/index.rst +++ b/docs/examples/pytorch/index.rst @@ -17,6 +17,7 @@ computer vision task with WeightsLab fully wired in. :maxdepth: 1 classification + fashion_mnist_umap segmentation detection clustering diff --git a/docs/examples/usecases/index.rst b/docs/examples/usecases/index.rst index 8155053d..662e2d49 100644 --- a/docs/examples/usecases/index.rst +++ b/docs/examples/usecases/index.rst @@ -4,7 +4,8 @@ Specific User Usecases ====================== Task-specific integrations that go beyond the standard loop: point-cloud -inputs, per-sample loss trajectory analysis, and per-layer training dynamics. +inputs, per-sample loss trajectory analysis, per-layer training dynamics, and +projections of the learned representation. .. raw:: html @@ -19,3 +20,4 @@ inputs, per-sample loss trajectory analysis, and per-layer training dynamics. lidar_detection loss_shape_classification model_signals + projection_fashion_mnist diff --git a/docs/examples/usecases/projection_fashion_mnist.rst b/docs/examples/usecases/projection_fashion_mnist.rst new file mode 100644 index 00000000..7e7add6f --- /dev/null +++ b/docs/examples/usecases/projection_fashion_mnist.rst @@ -0,0 +1,95 @@ +Projections on Fashion-MNIST +============================ + +.. raw:: html + +
+ Usecase + projection + umap + t-sne + notebook + beta +
+ +Runs on the same small CNN, showing the ways a projection reaches the +:ref:`studio-projection-board` (beta): the **built-in** live UMAP, and **your +own** algorithm, plugged in from a script or from the Studio notebook. +Background and every option: :doc:`../../projection`. + +Built-in live UMAP, while it trains +----------------------------------- + +**Example:** ``weightslab/examples/Usecases/wl-fashion-mnist-umap`` + +A plain training loop with nothing projection-specific in it. Wrapping the model +attaches the projection, because it is on by default; the dict only tunes it: + +.. code-block:: python + + model = wl.watch_or_edit( + FashionCNN().to(device), flag="model", device=device, + projection={"every_n_steps": 25, "graph_size": 1024}, # optional + ) + +Run ``python main.py``, press **Play** in the Studio, then open *Data +Exploration → Projection*. Until the first fit the board shows an orange *No +projection found* ribbon saying it is waiting; then one point per image appears, +coloured by class. The cloud starts as one tangle and pulls apart into clusters +as the representation forms — and the classes that stay mixed (Shirt, T-shirt, +Pullover, Coat) are the ones the model confuses. Lasso a mixed region to load +those images into the grid. + +The projection reads the input of the last ``Linear`` (``fc2``): the 128-D +representation the classifier decides from, not the logits. Test images are +placed by the encoder but never fitted on. + +Compare it with a PCA of your own +--------------------------------- + +**Example:** ``weightslab/examples/PyTorch/wl-fashion-mnist-umap`` + +Same CNN, driven in stages: train 1000 steps with the live UMAP, **reload** that +checkpoint, **discard** the 30 % of training samples with the lowest loss (the +ones the model already gets right), lay what is left out with a **2-D PCA you +write yourself**, then compare it with the UMAP in the board's picker and in +numbers (how often a point's nearest neighbours in the picture share its class). +A last stage retrains on what is left: the live UMAP follows the model, the PCA +stays the snapshot you took. The full walkthrough is +:doc:`../pytorch/fashion_mnist_umap`. + +Converge, pause, your own t-SNE +------------------------------- + +**Example:** ``weightslab/examples/Usecases/wl-fashion-mnist-custom-projection`` + +1. The built-in projection is **off** (``projection=False``): the projection in + this run is yours. +2. The script evaluates every 250 steps and, once test accuracy stops improving + (3 evaluations without +0.25 points), pauses itself: + + .. code-block:: python + + if plateau.update(test_accuracy): + wl.pause_training() # the next step waits for Play + +3. It has written ``own-projection.ipynb`` into the run directory, so the + Studio's notebook opens on it. Its cells collect the converged model's + features, run an exact **t-SNE written from scratch in plain torch**, and + hand the result over: + + .. code-block:: python + + wl.save_projection_coords(tsne_coords, batch_ids=ids, prefix="tsne") + +4. *Data Exploration → Projection* now offers ``tsne`` — and ``pca``, from the + notebook's last cell — in the picker. + +Running a notebook cell pauses training first and leaves it paused (see +:ref:`notebook-pauses-training`), so even a cell run before convergence sees one +state of the model. Press **Play** to resume training when you are done. + +Any algorithm plugs in the same way: anything that returns ``(N, 2)`` or +``(N, 3)`` coordinates for ``N`` sample ids, from ``sklearn``, ``openTSNE``, +``umap-learn`` or your own code. ``wl.project_dataset(model, loader, +method=TSNE(n_components=3))`` collects the features for you instead. diff --git a/docs/experiment_reports.rst b/docs/experiment_reports.rst index 196cedba..c0303f9f 100644 --- a/docs/experiment_reports.rst +++ b/docs/experiment_reports.rst @@ -21,7 +21,7 @@ Experiment Reports Ask the AI agent how your experiment is doing and it can produce a self-contained HTML report: signal trajectory plots, an automatic health classification per signal, dataset stats (sample counts, discard rate, tag -distribution), and a written analysis grounded in those exact numbers — +distribution), and a written analysis grounded in those exact numbers, branded with the WeightsLab logo. @@ -143,7 +143,7 @@ persist across separate sessions. Python/CLI callers that want the same overwrite-in-place behavior can pass the previous run's own path back in as ``output_path`` (:func:`ai_report_generation`) / ``--output`` (the CLI's ``report`` command) -— they already have direct control over the file, so there's no separate +, they already have direct control over the file, so there's no separate "update" flag for them. - **Weights Studio button**: the bar-chart icon immediately left of the @@ -216,7 +216,7 @@ Light / dark mode The report follows the browser's ``prefers-color-scheme`` automatically, and also has its own toggle button (top-right of the banner) for overriding that -— the choice is remembered (via ``localStorage``, scoped to that report file) +, the choice is remembered (via ``localStorage``, scoped to that report file) so reopening the same report keeps the theme you picked. Signal/distribution plots are rendered once by matplotlib on a fixed white canvas, so they sit in a small always-light thumbnail card in either theme, this keeps their own diff --git a/docs/index.rst b/docs/index.rst index 570323bf..c45588c3 100644 --- a/docs/index.rst +++ b/docs/index.rst @@ -12,7 +12,6 @@ Weightslab is a Python SDK to inspect, monitor, and edit training behavior for c
Quickstart Examples - Sandbox
@@ -49,6 +48,12 @@ Weightslab is a Python SDK to inspect, monitor, and edit training behavior for c Understand the 4-level workflow and each part independently. + .. grid-item-card:: Live 3-D Projection + :link: projection + :link-type: doc + + Watch the representation your model is learning, and select samples in it. + .. grid-item-card:: Notebooks :link: notebooks :link-type: doc @@ -120,6 +125,7 @@ Weightslab is a Python SDK to inspect, monitor, and edit training behavior for c four_way_approach signal_trajectory_classification + projection custom_evaluation diff --git a/docs/model_interaction.rst b/docs/model_interaction.rst index e4eeaa1b..b1d3f01e 100644 --- a/docs/model_interaction.rst +++ b/docs/model_interaction.rst @@ -14,7 +14,7 @@ Model wrapping parameters (``flag="model"``) -------------------------------------------- - Observe training signals at batch/sample granularity. -- Watch the model's own training dynamics, gradients, weights, activations — +- Watch the model's own training dynamics, gradients, weights, activations, per layer and per step (see `Training-dynamics signals`_). - Keep a stable ledger/proxy handle across runtime updates. - Enable dynamic controls without rewriting your loop architecture. diff --git a/docs/projection.rst b/docs/projection.rst new file mode 100644 index 00000000..b212c3bf --- /dev/null +++ b/docs/projection.rst @@ -0,0 +1,465 @@ +Live 3-D Projection (Parametric UMAP) +===================================== + +.. note:: + + **Beta.** The projection and its board work end to end, but their options + and on-disk format may still change between releases. + +A small encoder is trained alongside your model, on the UMAP objective, mapping +the model's penultimate features to 3-D. Each sample gets coordinates, written +back as ordinary per-sample signals, and Weights Studio draws them as a +navigable cloud (see :ref:`studio-projection-board`). + +It is **on by default and needs no code**. It never perturbs training: features +are detached at the hook and the encoder carries its own optimizer, so no +gradient from the projection can reach your model's weights. + +You can also bring **your own** projection , t-SNE, ``umap-learn``, PCA, +anything , and explore it on the same board (see +:ref:`projection-bring-your-own`). + +Quickstart +----------- + +Wrapping the model attaches it: + +.. code-block:: python + + import weightslab as wl + + model = wl.watch_or_edit(MyNet(), flag="model") + crit = wl.watch_or_edit(nn.CrossEntropyLoss(reduction="none"), + flag="loss", signal_name="loss_sample", + per_sample=True, log=True) + +After the first few fits, ``signals//umap_x``, ``_y`` and ``_z`` appear as +per-sample columns , so sorting, filtering, histograms and export all work on +them like any other signal. + +Turning it off +--------------- + +For the whole process:: + + WEIGHTSLAB_PROJECTION=0 python train.py + +``0``, ``false``, ``no``, ``off`` (any case) all disable it. Disabled means +*removed*: no hook, no encoder, no cost. + +For one model, from the model wrapper: + +.. code-block:: python + + model = wl.watch_or_edit(MyNet(), flag="model", projection=False) + +``projection=False`` also removes the hook of a model wrapped earlier in the +same process , there is one projection per process, and it follows the model +you wrapped last. The environment variable wins over everything: with +``WEIGHTSLAB_PROJECTION=0`` even ``projection=True`` or a ``projection={...}`` +dict installs nothing. + +Configuration +-------------- + +Environment variables (see also :doc:`configuration`): + +====================================== ========= ================================================== +Variable Default Meaning +====================================== ========= ================================================== +``WEIGHTSLAB_PROJECTION`` on ``0``/``false``/``no``/``off`` removes the feature. +``WEIGHTSLAB_PROJECTION_EVERY`` ``50`` Fit + write every N training steps. +``WEIGHTSLAB_PROJECTION_DIM`` ``3`` Output dimensions (2 draws on the z = 0 plane). +``WEIGHTSLAB_PROJECTION_NEIGHBORS`` ``15`` UMAP ``n_neighbors``. +``WEIGHTSLAB_PROJECTION_GRAPH`` ``512`` Samples the kNN graph is built over (see below). +``WEIGHTSLAB_PROJECTION_MAX_POINTS`` ``70000`` Points the board draws per view, at most 400,000. +====================================== ========= ================================================== + +Per model, through the ``projection`` keyword of +``wl.watch_or_edit(..., flag="model")``. A key set here takes precedence over +its environment variable: + +========================= =========== ===================================================== +Key Default Meaning +========================= =========== ===================================================== +``layer`` auto ``named_modules()`` name of the layer to hook (below). +``every_n_steps`` ``50`` As ``WEIGHTSLAB_PROJECTION_EVERY``. +``out_dim`` ``3`` As ``WEIGHTSLAB_PROJECTION_DIM``. +``n_neighbors`` ``15`` As ``WEIGHTSLAB_PROJECTION_NEIGHBORS``. +``graph_size`` ``512`` As ``WEIGHTSLAB_PROJECTION_GRAPH``. +``inner_steps`` ``4`` Encoder optimizer steps per fit. +``lr`` ``1e-3`` Encoder learning rate (Adam). +``min_dist`` ``0.1`` UMAP ``min_dist``: how tightly points may pack. +``spread`` ``1.0`` UMAP ``spread``: the scale of the embedding. +``repulsion_strength`` ``1.0`` Weight of the repulsive term. +``signal_prefix`` ``"umap"`` Column prefix: writes ``signals//_{x,y,z}``. +========================= =========== ===================================================== + +.. code-block:: python + + model = wl.watch_or_edit( + MyNet(), flag="model", + projection={"layer": "backbone.avgpool", "every_n_steps": 20, + "n_neighbors": 30, "min_dist": 0.05}, + ) + + # A string is shorthand for {"layer": ...}: + model = wl.watch_or_edit(MyNet(), flag="model", projection="backbone.avgpool") + +Why the graph is buffered +-------------------------- + +UMAP's graph is a **k-nearest-neighbour** graph, so it needs meaningfully more +samples than ``k``. Fitting on one training batch does not give it that: with a +batch of 16 and ``n_neighbors=15``, every sample is every other sample's +neighbour , the graph is complete, and there is no local structure left to +preserve. Measured on clustered features: + +========= ========================== =============================== +Batch Pairs that are edges Within/between membership ratio +========= ========================== =============================== +16 100 % 3 x +64 13 % 112 x +512 4 % ~10^11 x +========= ========================== =============================== + +So the features of **every** training batch go into a rolling buffer of the +most recent ``WEIGHTSLAB_PROJECTION_GRAPH`` samples, and each fit builds its +graph over that. Graph quality then no longer depends on whatever batch size +your training loop happens to use. **If your projection shows one blob instead +of clusters, this is the first thing to check** , raise +``WEIGHTSLAB_PROJECTION_GRAPH``, or lower ``WEIGHTSLAB_PROJECTION_NEIGHBORS``. + +Between fits, each training batch is also **placed** with the encoder as it +stands , one small forward pass, no gradient , so every sample of an epoch gets +coordinates from a recent encoder, not only the batches that happen to land on +a fit step. Nothing is placed before the first fit, so the board never shows a +cloud drawn by an untrained encoder. + +Both matter more than they sound. When only fit-step batches were buffered and +placed, the bundled ``wl-parametric-umap`` demo left a quarter of its samples +with no coordinates at all and separated its clusters 1–3x; with every batch +buffered and placed, all samples are drawn and the clusters separate ~25x. + +Which layer gets hooked +------------------------ + +One rule: **the input of the last parameterised layer** , last ``nn.Linear``, +else the last conv, else the last leaf module's output. That layer is the +model's head, so what flows *into* it is the representation; what comes *out* +is class scores. Auxiliary heads (``aux_classifier``, ``aux_head``, ...) are +skipped, and so are frozen layers while a trainable one exists: a layer whose +weights never train is a fixed operator, not the head (YOLO's ``dfl`` box +decoder, registered after the class/box convs, only runs at inference). + +The board names the layer under its title (``layer model.22 · C3k2 output``), +so a cloud never has to be taken on faith. ``WLAwareTrainer`` (Ultralytics) +hooks the output of the last block before the head, the layer UL's own +``YOLO.embed()`` pools, and also feeds the projection from the EMA model that +validation runs on. + +Reading the output instead fails silently: a segmentation head emits +``(B, num_classes)``, a well-shaped tensor UMAP will happily lay out, and you +would be looking at class scores believing they were the representation +(``fcn_resnet50`` gives ``(B, 512)`` features this way, ``(B, 21)`` logits the +other). + +The activation is reduced to ``(B, F)`` along the right axis: a Linear's input +has features last (so ``(B, T, C)`` pools over tokens and keeps ``C``), a conv's +input is channels-first. Tuple returns and HuggingFace output dicts are +unwrapped automatically. + +Override with ``projection={"layer": ""}``. Worth doing +when the model has multiple heads ("last Linear" means last *defined*, not last +executed) or when the auto-picked layer is not on the forward path , in which +case ``wl.projection.get_tracker().stats()`` reports ``fits: 0`` and a warning +is logged. An unknown name is not fatal: it is logged and the auto-pick is used. + +Verified across the bundled model zoo (MLPs, VGG, ResNet, U-Net/3+/3D, +TinyYOLO, FCN-ResNet50), torchvision, and transformers (GPT-2, BERT, +``BertForSequenceClassification``) , see +``tests/general/test_projection_baseline_models.py``. + +.. _projection-bring-your-own: + +Bring your own projection +-------------------------- + +The board draws **any** per-sample ``signals//_x`` / ``_y`` (and +optionally ``_z``) columns, whoever wrote them, and lists every prefix it finds +in its picker. So your own algorithm plugs in by writing coordinates under a +prefix of your choosing , beside the built-in ``umap``, or instead of it. + +**You have the coordinates** , ``wl.save_projection_coords``: + +.. code-block:: python + + from sklearn.manifold import TSNE + + # features: (N, F) array you collected; ids: the N sample ids (the ids the + # WeightsLab loaders yield, in the same order) + coords = TSNE(n_components=3).fit_transform(features) # (N, 3) + wl.save_projection_coords(coords, batch_ids=ids, prefix="tsne") + +**WeightsLab collects the features, your algorithm lays them out**, +``wl.project_dataset(..., method=...)``. ``method`` is anything with a +scikit-learn style ``fit_transform`` (``TSNE``, ``PCA``, ``umap.UMAP``) or a +plain callable ``(N, F) ndarray -> (N, 2 | 3)``. The features come from the same +layer the live projection uses, or from ``layer=`` if you pass one: + +.. code-block:: python + + import umap # pip install umap-learn + + wl.project_dataset(model, eval_loader, method=umap.UMAP(n_components=3), + prefix="umap_learn") + wl.project_dataset(model, eval_loader, method=TSNE(n_components=2), + layer="backbone.layer3") # prefix defaults to "tsne" + +Either call works whenever you like , after training, once per epoch, or at a +given step from inside your loop: + +.. code-block:: python + + for x, ids, y in train_loader: + with wl.guard_training_context: + ... # your training step + if model.get_age() == 500: + wl.project_dataset(model, eval_loader, method=TSNE(n_components=3), + prefix="tsne_step500", max_samples=20_000) + +It runs synchronously, so training waits for it: t-SNE on tens of thousands of +samples takes minutes, which is what ``max_samples`` is for. The model is used +in ``eval()`` under ``no_grad`` and its train/eval mode is restored afterwards. + +The sweep is a **whole pass, whatever training is doing**: + +- A tracked loader your training loop is mid-epoch on is restarted first (the + loop begins a fresh epoch afterwards), so ``train_loader`` is as good a choice + as an evaluation loader. Discarded samples are skipped by the loader itself, + as in training. +- It **does not age the model**: ``model.get_age()`` is the same before and + after, and the model's tracking mode is put back. +- The returned dict names the samples placed (``sample_ids``), in the order of + the feature rows your ``method`` received, so you can pair them with labels. + +The rules: + +- **Prefix** , letters, digits, ``_``, ``.``, ``-``, starting with a letter or + digit. Writing the same prefix again replaces the coordinates of the samples + you pass; samples you leave out keep theirs. Use a new prefix (such as + ``tsne_step500``) to keep a snapshot. +- **Not** ``umap`` **while the built-in runs** , it rewrites its prefix on every + fit and would overwrite yours sample by sample. Either name, or turn the + built-in off with ``projection=False`` and use the name freely. +- **2-D or 3-D** , ``(N, 2)`` is drawn on the z = 0 plane. Anything else is + refused with an error naming the shape. Rows of NaN mean "not placed" and are + skipped by the board. +- **Not counted as seen** , writing coordinates leaves the samples' + ``nb_seen`` / ``last_seen`` alone; placing a sample in a picture is not the + model training on it. +- **Mistakes fail fast** , a bad prefix, the live prefix, or an unknown + ``method`` string is refused *before* ``project_dataset`` sweeps the dataset. +- No encoder is saved or adopted for your own method; ``epochs``, + ``fit_batch``, ``lr``, ``min_dist``, ``spread``, ``n_neighbors`` and ``adopt`` + apply to the built-in only. + +When the built-in is off and only your projection exists, the board opens on it +directly; with both, ``umap`` comes first in the picker. The status line's fit +count belongs to the built-in, so it is shown only for its prefix. Until a +projection exists the board shows an orange *No projection found* ribbon with +the reason. + +**From the Studio notebook.** The notebook runs in the training process, so a +cell can read the live model and call ``wl.save_projection_coords`` directly. +Running a cell pauses training first (see :ref:`notebook-pauses-training`), so +the features you project come from one state of the model. To pause at a +moment of your choosing from the training script itself , say, once the run +has converged , call ``wl.pause_training()``; training waits at its next step +until Play. + +Three runnable examples: + +- ``examples/Usecases/wl-fashion-mnist-umap`` , Fashion-MNIST with the built-in + live UMAP running alongside training. +- ``examples/PyTorch/wl-fashion-mnist-umap`` , trains 1000 steps with the live + UMAP, reloads that checkpoint, discards the lowest-loss training samples, + projects what is left with a PCA you write yourself, and compares the two. See + :doc:`examples/pytorch/fashion_mnist_umap`. +- ``examples/Usecases/wl-fashion-mnist-custom-projection`` , trains to + convergence, pauses itself with ``wl.pause_training()``, and drops a notebook + into the run directory that computes a t-SNE written from scratch and plugs it + into the board. + +Checkpoints and restarts +------------------------- + +The encoder follows the model weights, so the cloud always describes the model +you actually have: + +- **Every model checkpoint** saves the encoder beside it, in a ``projection/`` + subdirectory next to the weight files + (``.../models//projection/.pt``). Never *inside* the + checkpoint: the encoder is not part of your model and must not ride along into + a deployable weights file. Never *beside* it either, where it would match the + weight-checkpoint glob. +- A **run-level copy** is kept at ``/projection/_encoder.pt`` + and refreshed every 20 fits and at every checkpoint, so a crash or a Ctrl-C + does not cost the layout. + +On a restore: + +.. list-table:: + :header-rows: 1 + :widths: 40 60 + + * - What happens + - Encoder you get + * - Restart; the wrapper auto-loads the latest checkpoint + - The encoder saved **with that checkpoint** , not the run-level copy, + which may be from later steps. + * - Restore a checkpoint from Studio or the agent + - The encoder saved with it; buffered features are dropped. + * - The checkpoint predates the first fit + - None: the layout starts again from that point. + * - The checkpoint restores a different architecture + - Its encoder, re-hooked onto the same layer (by name) of the restored + model. + * - No checkpoint restored (fresh weights) + - The run-level copy, so the layout continues. + +The coordinates follow the same rule as every other signal. A restore from +Studio or the agent rewinds per-sample state to the checkpoint's step: samples +seen after it lose their coordinates , your own projections' too , and are +placed again as training revisits them, so the cloud thins and refills. A plain +restart keeps the stored coordinates until the next fits overwrite them. + +``wl.projection.save_projection()`` / ``load_projection()`` save and load the +run-level copy explicitly. + +When the projection cannot run +------------------------------- + +The projection is instrumentation, so the contract is simple: **it never raises +into your training loop and never stops a run.** What you get instead: + +.. list-table:: + :header-rows: 1 + :widths: 40 60 + + * - Situation + - Behaviour + * - It cannot be installed (no hookable layer, an unknown option) + - A warning; the model is wrapped normally, with no projection. + * - The hooked layer is not on the forward path + - A warning after 3 visits naming the fix (``projection={"layer": ...}``). + * - A fit raises (out of memory, an unexpected shape) + - A warning the first time; it retries at the next fit. + * - 5 fits in a row fail + - Turned off for the rest of the run, with a warning saying why + (``stats()["disabled_reason"]``). Coordinates already written stay on + the board. + * - The model's features are NaN/inf (it diverged) + - Those batches are skipped, warned once, so the encoder keeps its last + good layout instead of turning NaN for good. + * - The UMAP loss is NaN/inf + - The step is not taken, so the encoder's weights stay finite; it counts + as a failed fit. + * - Coordinates cannot be written + - A warning once; training is unaffected. + * - The hook itself errors + - Ignored; the forward pass is never affected. + +The explicit calls are different: ``wl.save_projection_coords`` and +``wl.project_dataset`` are things you invoke yourself, so bad input raises there +(``ValueError`` / ``TypeError`` with what was expected). + +Re-projecting offline +---------------------- + +When a run ends and the live projection watched the wrong layer, you do not +have to retrain: + +.. code-block:: python + + model = wl.watch_or_edit(MyNet(), flag="model") # checkpoint restores + loader = wl.watch_or_edit(train_ds, flag="data", loader_name="train_loader") + + wl.project_dataset(model, loader, layer="backbone.layer3", + prefix="umap_layer3", epochs=20) + +No training step runs , the model is used in ``eval()`` under ``no_grad``, and +its weights and train/eval mode are left untouched. It also fits better than +the live path: every feature is collected first, so each epoch draws +neighbourhoods from the whole dataset rather than one batch. + +Different ``prefix`` values coexist as separate columns, so several candidate +layers can be compared; the board's picker switches between them. Pass +``method=`` to use your own algorithm instead (see +:ref:`projection-bring-your-own`). + +Resuming training afterwards: with a **different** prefix the offline result is +a frozen snapshot and the live projection carries on independently. With the +**same** prefix, the live projection *adopts* the offline fit by default +(``adopt=None``) , it re-hooks to that layer and keeps refining it , because +not adopting would let training overwrite those coordinates with a fresh +encoder's output, mixing two layouts. Force it with ``adopt=True``/``False``. + +Progress is reported with ``tqdm`` (``features`` / ``fit`` / ``encode``); +``verbose=False`` silences it. Returns a dict with ``samples``, +``feature_dim``, ``prefix``, ``columns``, ``method``, ``final_loss`` and +``adopted_by_live`` , check ``feature_dim`` to confirm you hooked what you +meant to. + +Task coverage +-------------- + +The projection needs only a forward pass through the hooked layer and a +per-sample write carrying ``batch_ids`` , either a ``flag="loss"`` criterion or +a direct :func:`wl.save_signals`. Every supervised usecase provides one. + +Detection/segmentation with ``per_instance=True`` is handled: the projection +receives the **sample**-level ids, not the per-annotation ones, so there is one +point per image. Sequence models pool over tokens, giving one point per +sequence. + +Train and eval splits both appear. Eval batches are **placed but never fit on** +, held-out samples belong in the picture, but letting them shape the layout +would leak the eval set into it. + +Debugging +---------- + +``wl.projection.get_tracker()`` is ``None`` when no projection is attached; +otherwise ``.stats()`` reports ``enabled``, ``layer``, ``feature_dim``, +``fits``, ``rebuilds``, ``buffered``, ``samples_written``, ``last_loss``, +``failures``, ``skipped_nonfinite``, ``disabled_reason`` and the ``columns`` it +writes. ``WEIGHTSLAB_LOG_LEVEL=DEBUG`` shows every retried fit. + +Limits +------- + +- The layout is **not comparable across runs**, nor across an architecture + edit: the encoder is rebuilt when the hooked layer's width changes + (``stats()["rebuilds"]`` counts it) and the cloud re-lays-out. Compare + structure, not coordinates. +- **Token pooling ignores the attention mask** , a forward hook cannot see it, + so a padded language batch averages padding positions in. Harmless at uniform + lengths; with varying lengths, hook your own pooling module instead. +- Coordinates exist only for samples seen since the projection attached. +- A fit needs at least 4 buffered samples, and ``n_neighbors`` is clamped to the + graph size minus one. +- The picker lists only prefixes a WeightsLab writer registered (the live + projection, ``wl.project_dataset``, ``wl.save_projection_coords``), recorded + in ``/projection/prefixes.json``. Columns you write yourself + with ``wl.save_signals`` as ``tsne_x`` / ``tsne_y`` are *not* offered; use + ``wl.save_projection_coords``. A run with no such file (written before it + existed) still lists every ``_x``/``_y`` column pair. + +See also +--------- + +- :ref:`studio-projection-board` , the Projection Board in Weights Studio. +- :doc:`signal_trajectory_classification` , the other "look at the shape, not + the value" view, over time rather than over the representation. +- :doc:`configuration` , every environment variable in one table. diff --git a/docs/quickstart.rst b/docs/quickstart.rst index 8388f824..61f812da 100644 --- a/docs/quickstart.rst +++ b/docs/quickstart.rst @@ -176,10 +176,10 @@ Migrating a real codebase? The full guides are step += 1 Three things to notice. The tracked ``train_loader`` yields - ``(inputs, uids, labels)`` — those ``uids`` are what tie a loss value back + ``(inputs, uids, labels)`` , those ``uids`` are what tie a loss value back to the sample that produced it, which is why they are handed to the loss as ``batch_ids``. The loop is open-ended: you stop it from the studio, not - with a step budget. And it does not start on its own — run + with a step budget. And it does not start on its own , run ``weightslab start`` in another terminal and press **Play**. .. tab-item:: WeightsLab From TensorBoard @@ -414,7 +414,7 @@ Migrating a real codebase? The full guides are One thing has no equivalent, and is not meant to: there are no sweeps. - WeightsLab is built around staying inside *one* run and steering it — + WeightsLab is built around staying inside *one* run and steering it, raise the learning rate when the curve flattens, discard the samples poisoning it, keep going. If you need a sweep, keep the tool you sweep with. diff --git a/docs/resource_monitoring.rst b/docs/resource_monitoring.rst index d262d5dd..985aa5b6 100644 --- a/docs/resource_monitoring.rst +++ b/docs/resource_monitoring.rst @@ -82,7 +82,7 @@ To turn everything off, either: export WEIGHTSLAB_DISABLE_RESOURCE_MONITORING=1 -or set ``enabled: false`` in ``resource_monitoring.yaml`` (see below) — +or set ``enabled: false`` in ``resource_monitoring.yaml`` (see below), the YAML value wins if both are set. Enabling only specific categories @@ -97,7 +97,7 @@ Two ways to restrict which categories are sampled: export WL_RESOURCE_MONITOR_CATEGORIES=cpu,memory,gpu -- **YAML file** (``resource_monitoring.yaml``), per-category booleans — +- **YAML file** (``resource_monitoring.yaml``), per-category booleans, lets you leave everything on and disable just one or two: .. code-block:: yaml diff --git a/docs/segmentation_usecase.rst b/docs/segmentation_usecase.rst index ead642f8..ebbfb850 100644 --- a/docs/segmentation_usecase.rst +++ b/docs/segmentation_usecase.rst @@ -164,7 +164,7 @@ Important: --------------------------------------------------- ``utils/criterions.py`` also registers free-form signals via ``custom_signals()`` -— a static signal computed from the image, and a dynamic signal that reacts to a +, a static signal computed from the image, and a dynamic signal that reacts to a logged metric: .. code-block:: python diff --git a/docs/signal_trajectory_classification.rst b/docs/signal_trajectory_classification.rst index 0ae8c432..fd093d47 100644 --- a/docs/signal_trajectory_classification.rst +++ b/docs/signal_trajectory_classification.rst @@ -103,7 +103,7 @@ own vocabulary; yours can return anything. - ``@wl.signal_classifier(signal="loss_sample")``, classify only that one signal. -- ``@wl.signal_classifier`` / ``@wl.signal_classifier()`` (no ``signal=``) — +- ``@wl.signal_classifier`` / ``@wl.signal_classifier()`` (no ``signal=``), become the global default for every signal that doesn't have its own per-signal classifier registered. diff --git a/docs/user_commands.rst b/docs/user_commands.rst index fdc597fc..2c190939 100644 --- a/docs/user_commands.rst +++ b/docs/user_commands.rst @@ -269,7 +269,7 @@ consequences: - ``--listen-port``, ``-p`` *(int)*, local port to expose. Default: **50051** (the port ``weightslab start`` proxies to by default, leave it unless you pass ``--backend-port`` or set ``GRPC_BACKEND_PORT``). -- ``--listen-host`` *(str)*, interface to bind. Default: **auto** — +- ``--listen-host`` *(str)*, interface to bind. Default: **auto**, ``127.0.0.1`` on Windows/macOS, ``0.0.0.0`` (all interfaces) on Linux. With the UI on the same machine, ``--listen-host 127.0.0.1`` works on Linux too and keeps the tunnel private. diff --git a/docs/user_functions.rst b/docs/user_functions.rst index a6440b30..f2eb5e89 100644 --- a/docs/user_functions.rst +++ b/docs/user_functions.rst @@ -16,9 +16,15 @@ Core registration and serving: - ``wl.watch_or_edit`` - ``wl.guard_training_context`` / ``wl.guard_testing_context`` - ``wl.start_training`` +- ``wl.pause_training`` - ``wl.serve`` - ``wl.keep_serving`` +Projection *(beta, see* :doc:`projection` *)*: + +- ``wl.project_dataset`` *(offline re-projection; built-in UMAP or your own method)* +- ``wl.save_projection_coords`` *(plug in coordinates you computed yourself)* + Signals: - ``wl.signal`` *(decorator, custom static/dynamic signals)* @@ -211,6 +217,59 @@ your training loop, optionally blocking first. for step, batch in enumerate(train_loader): ... +pause_training +-------------- + +**Signature** + +.. code-block:: python + + wl.pause_training() + +**Purpose** + +Pause training from your own code: the Studio's pause button, as a call, and +the counterpart of ``start_training``. The step in flight is not interrupted; +the next ``with wl.guard_training_context:`` blocks until training is resumed +(Play in the Studio header, or ``wl.start_training()``). While paused, the +Studio, the gRPC server and the notebook keep running against the live model. + +**Typical usage** + +.. code-block:: python + + for step, batch in enumerate(train_loader): + with wl.guard_training_context: + ... + if converged(): + wl.pause_training() # inspect the converged run, then press Play + +Running a cell in the Studio notebook pauses training the same way (see +:ref:`notebook-pauses-training`). + +project_dataset / save_projection_coords +---------------------------------------- + +**Signatures** + +.. code-block:: python + + wl.project_dataset(model, dataloader, layer=None, prefix=None, method="umap", ...) + wl.save_projection_coords(coords, batch_ids, prefix) + +**Purpose** + +Put a 2-D or 3-D projection of your samples on the Projection Board. +``project_dataset`` collects the model's features from one layer and lays them +out, with the built-in parametric UMAP or with your own ``method`` (anything +with ``fit_transform``, or a callable). ``save_projection_coords`` stores +coordinates you computed entirely yourself. Both write +``signals//_{x,y,z}`` per sample. ``project_dataset`` sweeps the whole +loader (restarting a tracked loader a training loop is mid-epoch on), does not +age the model, and returns a dict that includes ``sample_ids``, the samples +placed in the order of the feature rows your ``method`` received. See +:doc:`projection`. + serve ----- @@ -631,7 +690,7 @@ trajectory_stats **Purpose** -Scale- and noise-invariant summary statistics of one sample's value trajectory — +Scale- and noise-invariant summary statistics of one sample's value trajectory, the reusable feature layer :func:`classify_loss_shape` is built on. Returns ``None`` when *values* has fewer than 2 points. Build a custom :func:`signal_classifier` on top instead of re-deriving these features by hand. @@ -1975,7 +2034,7 @@ export_annotations **Purpose** -Export bounding-box/segmentation annotations to a relabeling-tool format — +Export bounding-box/segmentation annotations to a relabeling-tool format, the Python-API counterpart to Weights Studio's "Export" button and the ``weightslab export`` CLI command. See :doc:`export` for the full format reference and known limitations (image-path/class-name resolution). diff --git a/docs/weights_studio/notebook.rst b/docs/weights_studio/notebook.rst index 4aaadc7a..e8a79ab8 100644 --- a/docs/weights_studio/notebook.rst +++ b/docs/weights_studio/notebook.rst @@ -31,6 +31,13 @@ How it works ``root_log_dir``. Reopening the panel, even after restarting the UI, as long as it points at the same experiment, reloads the same cells, their source, and their last-run outputs. +- A new experiment's notebook starts with two worked cells: the data and + logged signals at a glance, then the model's layers (by index and name, the + one the projection reads marked) with how to re-project from another layer. + Set ``LAYER`` in that cell to a name or index from the list and run it + again: it calls ``wl.project_dataset(..., layer=...)`` on the live prefix, + and training carries on from that fit. To keep the layer for the next run, + pass ``projection={"layer": ""}`` to ``wl.watch_or_edit``. - Every cell runs against the training process's ONE shared kernel: only one cell executes at a time. Clicking Run on a second cell while another is still running queues it rather than firing a second concurrent execution. @@ -38,6 +45,9 @@ How it works (merged, in order), the value of the last expression, any ``matplotlib`` figures rendered inline as images, and a full traceback on error. - A run can be interrupted mid-flight with the stop button next to the cell. +- Running a code cell pauses a running training first, and training stays + paused until Play: a cell always sees one state of the model (see + :ref:`notebook-pauses-training`). Cell types ---------- diff --git a/docs/weights_studio_ui/landing_page.rst b/docs/weights_studio_ui/landing_page.rst index c8beefe2..1bd75991 100644 --- a/docs/weights_studio_ui/landing_page.rst +++ b/docs/weights_studio_ui/landing_page.rst @@ -141,6 +141,37 @@ How it works (merged, in order), the value of the last expression, any ``matplotlib`` figures rendered inline as images, and a full traceback on error. - A run can be interrupted mid-flight with the stop button next to the cell. +- Running a code cell **pauses training first** (see below). + +.. _notebook-pauses-training: + +Running a cell pauses training +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +A cell runs in the training process, against the live model and data. If the +training loop kept stepping while a cell ran, a cell that reads features, +computes a t-SNE or inspects weights would see a model that changes underneath +it. So executing a code cell pauses a running training first — the same pause +as the header's button, so the header shows it paused — and waits for the +training step in flight to finish (up to 30 s) before the cell starts. + +Training then **stays paused**: resuming is your call, with Play. The first +cell that paused it says so in its output:: + + [WeightsLab] Training paused to run this cell -- press Play in the header to resume. + +Code sent to the embedded kernel some other way (``jupyter console +--existing``) pauses training the same way. Generating code with a ``>`` cell +does not pause anything; running the generated code does. + +To pause at a moment of your choosing from the training script instead — once +the run has converged, say — call ``wl.pause_training()``. + +Which notebook opens: the one used last in this session, else the most +recently written ``.ipynb`` under ``root_log_dir``, else a new +``notebook.ipynb``. A script can therefore ship a ready-made notebook by writing +it into ``root_log_dir`` before ``wl.serve()`` (see +``examples/Usecases/wl-fashion-mnist-custom-projection``). Cell types ~~~~~~~~~~~ diff --git a/docs/weights_studio_ui/main_area.rst b/docs/weights_studio_ui/main_area.rst index cec10734..83e31ab0 100644 --- a/docs/weights_studio_ui/main_area.rst +++ b/docs/weights_studio_ui/main_area.rst @@ -361,3 +361,172 @@ frame slider, so you can land on the exact frame a signal spiked on. **Volumetric images** get a Z-slice slider, and **text samples** render as text rather than as an image. + +.. _studio-projection-board: + +Projection Board +---------------- + +.. figure:: ../_static/screenshots/projection-board.png + :alt: Projection board showing the 3-D parametric-UMAP cloud + :width: 100% + +.. note:: + + **Beta** — the board carries a *Beta* pill in its header. It works end to + end, but its controls may still change between releases. + +A navigable 3-D view of the representation the model is learning: one point per +sample, placed by the live parametric-UMAP encoder (see :doc:`../projection` +for the training-side half), or by a projection you computed yourself. + +The board is opened from the **Projection** button in the Data Exploration +header, and opens beside the Data board. Nothing about a run puts it on screen +by itself. Opened before there is anything to draw — before the first fit, with +the projection turned off, or before you plugged in your own — it shows an +orange **No projection found** ribbon with the backend's reason, and fills in +as soon as coordinates arrive — no reload needed. + +Navigating +~~~~~~~~~~ + +- **Drag** to orbit the cloud. +- **ctrl+scroll** to zoom toward the cursor — the same gesture the plots use. + Plain scroll is left to the page, so the board never traps your wheel. +- **Fit** re-frames the whole cloud. +- The **expand** button gives the projection the whole view (Esc to exit). +- Keyboard: ``L`` lasso · ``R`` reset · ``C`` clear · ``F`` fit · ``E`` expand. + +Which projection +~~~~~~~~~~~~~~~~ + +A run can hold several projections: the live ``umap`` one, offline +re-projections from other layers, and any you computed yourself with t-SNE, +``umap-learn`` or anything else (``wl.save_projection_coords`` / +``wl.project_dataset(method=...)``, see :ref:`projection-bring-your-own`). When +there is more than one, a picker in the header switches between them; with the +built-in turned off, the board opens directly on yours. + +The status line counts points, e.g. +``4,812 pts (of 61,004 in view) · 70,000 projected · 120 fits``. The fit count +belongs to the live encoder, so it is shown only for its projection. + +The axes are unlabelled on purpose: UMAP coordinates carry no units, so they +are an orientation cue only. + +Selecting samples +~~~~~~~~~~~~~~~~~ + +Double-click a point, or switch on **Lasso** and drag a loop around a region. +Either way the result lands in the :ref:`studio-data-board` as an ordinary +filter, so the grid, the list view and every existing affordance work on it. A +single click only names the point (``sample 1234``). + +The lasso selects **every sample inside the loop, drawn or not**. The cloud on +screen is a level-of-detail subsample (see below), so the selection is made on +the server over the whole projection and applied to the Data board from there; +the status line shows its size. While the lasso is armed, a line over the cloud +says what the next loop will do: + +- **Shift** adds what is inside the loop to the current selection. +- **Ctrl** *refines* it: only the samples already selected that are also inside + the new loop are kept. A single loop selects a whole column through the + cloud; rotate the cloud and Ctrl-lasso the same cluster again, and the two + columns intersect in a 3-D volume — the cluster alone. The loop is drawn + dashed cyan while Ctrl is held. + +With **Overview** following a grid page, every point of the page is drawn, so +the lasso simply selects the drawn points. + +The projection itself is **never** filtered by that selection — it draws the +whole dataset, with the selection highlighted. Following its own filter would +collapse the cloud to the points just selected, leaving nothing to select from +next. **Clear** drops the selection and the highlight. + +**Overview** decides whether the cloud follows the Data board's *other* +filters (a query you or the agent applied). Off by default: the cloud shows the +whole dataset, which is what lets you see where a filtered subset sits in the +representation. Switch it on to draw only the Data board's current view. + +The other way round, right-click a sample in the grid and choose **Highlight in +Projection** to find it in the cloud. Right-clicking inside the cloud itself is +the same as **Reset view**. + +Colour +~~~~~~ + +The settings cog holds the colour picker, the point size and **Reset view** +(camera, axes, selection and highlight). + +Points are coloured by **Class** (the ground-truth label) by default — on a +classification set, the classes are what the clusters mean. **Split** uses the +grid's split palette, so a sample is the same colour everywhere in Studio. Any +metadata column works too: numeric columns get a ramp, categorical ones a +palette. Discarded samples are drawn grey, still in place. + +Level of detail +~~~~~~~~~~~~~~~ + +The board never downloads the whole dataset. It asks for the points its camera +can see — the view frustum — up to a render budget (70,000 by default), and +the server answers with that many of them; zooming in narrows the frustum, so +the same budget buys finer detail. Every camera move that settles asks again. +The subtitle says how many are drawn; hovering it says of how many in view. + +The budget is set where the projection is served, for every viewer at once: +``WEIGHTSLAB_PROJECTION_MAX_POINTS`` in the training process's environment, at +most 400,000. + +.. code-block:: bash + + WEIGHTSLAB_PROJECTION_MAX_POINTS=150000 python train.py + +It is read on every request, so in a run that is already going, setting it from +the Studio notebook applies at the next camera move: + +.. code-block:: python + + import os + os.environ["WEIGHTSLAB_PROJECTION_MAX_POINTS"] = "150000" + +One browser can override it for itself, in its developer console (remove the +key to go back to the server's value): + +.. code-block:: javascript + + localStorage.setItem('projection-render-budget-v1', '200000') + +Two properties follow from the sampling being deterministic rather than random: +repeated requests pick the same points (the cloud does not boil when you nudge +the camera), and zooming in only *adds* points. The sample is also +**stratified** by cluster, at a density proportional to ``1/sqrt(cluster +size)``: a cluster 100 times smaller is drawn 10 times denser than a flat +sample would draw it, so a small cluster cannot be sampled out of existence and +read as noise, and every class still appears when the dataset is many times the +budget. + +The server keeps an index of the projection — the sampling order and a spatial +grid — rebuilt in the background as training moves the points, so a view costs +about what it returns rather than the size of the dataset: about 0.1 s for +50,000 points whether the projection holds one million samples or ten. The +first open of a large projection builds that index once (a few seconds at 10M +samples). + +With **Overview** following a grid page, the cloud is exactly that page, drawn +whole: no budget and no frustum. + +The status line says plainly when the view is decimated (the ``of ... in +view`` part), so a level-of-detail subsample never reads as all the data there +is. + +Arranging the boards +-------------------- + +Every board in the main area — Plots, Data, Projection, and any added later — +is dragged by its header title: + +- drop on the **top or bottom half** of another board to stack above or below it; +- drop on a **left or right edge** to sit alongside it, splitting the width. + +The board you are dragging over highlights the exact edge you will land +against, and the arrangement is remembered between sessions. diff --git a/docs/whats_new.rst b/docs/whats_new.rst index 9761a516..04e274a0 100644 --- a/docs/whats_new.rst +++ b/docs/whats_new.rst @@ -207,7 +207,7 @@ list lives. > Re-opened from #265 with the Ultralytics example notebooks excluded (kept on ``dev`` only, not ready for this release). - **WeightsLab v1.4.0 —** + **WeightsLab v1.4.0** Release Notes diff --git a/pyproject.toml b/pyproject.toml index 256a093c..78ecd407 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -183,6 +183,9 @@ weightslab = [ # project config from the directory it is started in, which is that same # workspace; this one points `instructions` at AGENTS.md. "opencode.json", + # Last release tag, kept current by release.yml; the final fallback for + # weightslab.__version__ (see _resolve_version in __init__.py). + "VERSION", ] [tool.setuptools.exclude-package-data] @@ -203,3 +206,13 @@ local_scheme = "no-local-version" markers = [ "scale: large-history robustness tests (sweep-scale fixtures; deselect with -m 'not scale')", ] +filterwarnings = [ + # "tag:..." / dotted column names can never be Python identifiers. + # h5_dataframe_store silences this at import, but pytest resets warning + # filters after every test, so the suite needs its own copy. + "ignore::tables.NaturalNameWarning", + # Lightning's DataLoader performance hint; the tests use num_workers=0 on purpose. + "ignore:.*does not have many workers:UserWarning", + # ipykernel's own code (asyncio event-loop policies on Python 3.14+). + "ignore::DeprecationWarning:ipykernel", +] diff --git a/tests/backend/test_ledgers.py b/tests/backend/test_ledgers.py index 513b1caf..06c1e209 100644 --- a/tests/backend/test_ledgers.py +++ b/tests/backend/test_ledgers.py @@ -717,5 +717,26 @@ def reset(self): self.assertEqual(wrapped.calls, 1) +class LedgerClearDropsCachedHandlesTests(unittest.TestCase): + """ledgers.clear_all() must not leave src writing to the old dataframe. + + src caches the dataframe proxy in DATAFRAME_M. clear() drops that proxy and + the next registration makes a new one, so a kept handle sent every + save_signals write to the previous run while the new run's table stayed + empty -- what wl.clear_all() already prevented, and ledgers.clear_all() + (the call most tests and integrations make) did not. + """ + + def tearDown(self): + GLOBAL_LEDGER.clear() + + def test_clear_resets_src_cached_dataframe(self): + from weightslab import src + src.DATAFRAME_M = GLOBAL_LEDGER.get_dataframe(DEFAULT_NAME) + self.assertIsNotNone(src.DATAFRAME_M) + GLOBAL_LEDGER.clear() + self.assertIsNone(src.DATAFRAME_M) + + if __name__ == "__main__": unittest.main() diff --git a/tests/backend/test_overlapping_loader_loops.py b/tests/backend/test_overlapping_loader_loops.py new file mode 100644 index 00000000..ad131cc9 --- /dev/null +++ b/tests/backend/test_overlapping_loader_loops.py @@ -0,0 +1,65 @@ +"""Two `for batch in loader:` loops over the same loader at once. + +The training script's own test pass and an evaluation the Studio triggers can +iterate the same loader together, the first one paused mid-epoch. The loop +that finished first used to clear the loader's shared `is_a_loop` flag under +the other, whose epoch then never ended: the script's `for batch in +test_loader:` ran forever and training froze (the classification e2e's +"train for N steps" hang, after a test-split evaluation). +""" +import unittest + +import torch +from torch.utils.data import TensorDataset + +import weightslab.data.data_samples_with_ops as _dso +from weightslab.backend.dataloader_interface import DataLoaderInterface +from weightslab.backend.ledgers import Proxy + +N, BATCH = 10, 2 # five batches an epoch +CAP = 50 # "endless", for a loop that should stop after five + + +def drain(iterator): + count = 0 + for _ in iterator: + count += 1 + if count >= CAP: + break + return count + + +class TestOverlappingForLoops(unittest.TestCase): + def setUp(self): + with _dso._REGISTRY_LOCK: + _dso._GLOBAL_UID_REGISTRY.clear() + data = TensorDataset(torch.arange(N, dtype=torch.float32).unsqueeze(1), torch.zeros(N)) + self.loader = Proxy(DataLoaderInterface(data, batch_size=BATCH, shuffle=False)) + + def test_a_paused_loop_still_ends_after_another_one_finishes(self): + outer = iter(self.loader) # the script's test pass ... + next(outer) # ... paused mid-epoch + self.assertLess(drain(iter(self.loader)), CAP) # the evaluation, to its end + self.assertLess(drain(outer), CAP, "the paused loop never saw its epoch end") + + def test_without_a_for_loop_plain_next_rolls_over_epochs(self): + drain(iter(self.loader)) + outer = iter(self.loader) + next(outer) + drain(iter(self.loader)) + drain(outer) + self.assertFalse(self.loader.is_a_loop) + # The training loop's own `next(loader)` keeps going across epochs. + for _ in range(3 * N // BATCH): + next(self.loader) + + def test_an_abandoned_loop_releases_its_hold(self): + it = iter(self.loader) + next(it) + it.close() # what a managed evaluation does when it stops early + it.close() # once only + self.assertFalse(self.loader.is_a_loop) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/backend/test_version.py b/tests/backend/test_version.py new file mode 100644 index 00000000..28e29bf1 --- /dev/null +++ b/tests/backend/test_version.py @@ -0,0 +1,48 @@ +"""Tests for weightslab._resolve_version's committed-release fallback. + +weightslab/VERSION holds the last release tag (kept current by release.yml's +sync-version-file job) and is the source of __version__ — and so of the +telemetry version — only once git, _version.py and dist metadata all fail. +""" +import os +import re +import sys +import types +import unittest +from importlib.metadata import PackageNotFoundError +from unittest.mock import patch + +import weightslab + +_VERSION_FILE = os.path.join(os.path.dirname(weightslab.__file__), "VERSION") +# Normalized PEP 440 public version, as release.yml writes it (no "v", no local). +_PEP440 = re.compile(r"^\d+(\.\d+)*((a|b|rc)\d+)?(\.post\d+)?(\.dev\d+)?$") + + +def _read_version_file() -> str: + with open(_VERSION_FILE, encoding="utf-8") as fh: + return fh.read().strip() + + +class TestReleaseVersionFile(unittest.TestCase): + def test_version_file_is_a_normalized_version(self): + self.assertRegex(_read_version_file(), _PEP440) + + def test_used_when_git_build_and_metadata_all_fail(self): + # Not the main process -> the git step is skipped; a None module entry + # makes `from ._version import ...` raise ImportError. + with patch.object(weightslab, "_IS_MAIN_PROCESS", False), \ + patch.dict(sys.modules, {"weightslab._version": None}), \ + patch("importlib.metadata.version", side_effect=PackageNotFoundError): + self.assertEqual(weightslab._resolve_version(), _read_version_file()) + + def test_not_used_while_a_build_version_exists(self): + built = types.ModuleType("weightslab._version") + built.__version__ = "9.9.9.dev1" + with patch.object(weightslab, "_IS_MAIN_PROCESS", False), \ + patch.dict(sys.modules, {"weightslab._version": built}): + self.assertEqual(weightslab._resolve_version(), "9.9.9.dev1") + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/components/test_global_monitoring_unit.py b/tests/components/test_global_monitoring_unit.py index ee88c790..5c68534a 100644 --- a/tests/components/test_global_monitoring_unit.py +++ b/tests/components/test_global_monitoring_unit.py @@ -72,5 +72,44 @@ def test_guard_context_suppresses_runtime_error(self): self.assertTrue(suppressed) +class TestPauseControllerCheckpointManager(unittest.TestCase): + """resume() dumps through the manager registered NOW, not the first one. + + The controller is a process-wide singleton. It used to keep the manager it + first saw, so after ledgers.clear_all() and a new experiment it went on + writing the new run's checkpoints into the previous run's manifest, and a + reload by the current hash found nothing. + """ + + def tearDown(self): + from weightslab.backend import ledgers + ledgers.clear_all() + + def test_resume_follows_the_manager_registered_after_clear_all(self): + from unittest.mock import MagicMock + from weightslab.backend import ledgers + from weightslab.components.global_monitoring import PauseController + + first, second = MagicMock(name="first"), MagicMock(name="second") + first.hash_by_module = second.hash_by_module = ["aaaa", "bbbb", "cccc"] + controller = PauseController() + with patch("weightslab.components.global_monitoring.set_hyperparam"): + ledgers.register_checkpoint_manager(first) + controller.resume(force=True) + first.save_pending_changes.assert_called() + first.reset_mock() + + ledgers.clear_all() + ledgers.register_checkpoint_manager(second) + controller.resume(force=True) + + # Everything after the switch goes to the new run's manager, nothing to + # the old one (call counts are not pinned: other code may use it too). + second.update_experiment_hash.assert_called_with(first_time=True) + second.save_pending_changes.assert_called() + first.update_experiment_hash.assert_not_called() + first.save_pending_changes.assert_not_called() + + if __name__ == "__main__": unittest.main() diff --git a/tests/conftest.py b/tests/conftest.py index 30bcc6f0..69940bf8 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -68,3 +68,35 @@ def _disable_resource_monitor(): os.environ.pop("WEIGHTSLAB_DISABLE_RESOURCE_MONITORING", None) else: os.environ["WEIGHTSLAB_DISABLE_RESOURCE_MONITORING"] = previous + + +# The test modules that spawn (fake) OpenCode servers. +_OPENCODE_SPAWNING_MODULES = ( + "test_opencode_process", + "test_opencode_shared_server_integration", + "test_server_agent", +) + + +@pytest.fixture(autouse=True) +def _private_opencode_machine_lock(request, monkeypatch): + """Give each OpenCode-spawning test its own machine-wide lock. + + opencode_process publishes every server it spawns in + ``~/.weightslab/.wl_opencode.json`` (the machine-wide rendezvous), and the + fake servers these tests start are detached on purpose, so they outlive + the test. Against the real home directory, the next test found the + previous one's fake still healthy and adopted it instead of spawning + (``'machine-lockfile' != 'spawned'``), and running the suite overwrote the + developer's own lock with a test server's address. + + test_opencode_rendezvous isolates the home directory itself and asserts the + real path layout, so it is left alone here. + """ + if request.module.__name__.rsplit(".", 1)[-1] not in _OPENCODE_SPAWNING_MODULES: + yield + return + from weightslab import opencode_process + lock = request.getfixturevalue("tmp_path") / ".weightslab" / opencode_process.LOCK_FILENAME + monkeypatch.setattr(opencode_process, "machine_lock_path", lambda: lock) + yield diff --git a/tests/data/test_video_utils.py b/tests/data/test_video_utils.py index f9073396..6054822e 100644 --- a/tests/data/test_video_utils.py +++ b/tests/data/test_video_utils.py @@ -314,3 +314,136 @@ def test_encode_empty_clip_is_safe(): empty = np.zeros((0, 8, 8, 3), dtype=np.uint8) assert vu.encode_video_mp4(empty, fps=8) == b"" assert vu.encode_video_gif(empty, fps=8) == b"" + + +# --------------------------------------------------------------------------- +# Per-frame annotations (video segmentation / detection) +# --------------------------------------------------------------------------- +class TestSelectFrameAnnotation: + """A clip's annotation must be reduced to the frame the poster shows. + + The grid and the modal's still view draw exactly one frame. Sending the + whole [T, H, W] stack instead does not fail loudly — it gets RLE'd into a + stat the renderer reads as 2-D, i.e. a plausible-looking wrong overlay. + """ + + def test_poster_index_is_the_middle_frame(self): + # video_poster_frame() renders arr[T // 2]; anything drawn on it has + # to come from the same index or the overlay describes another moment. + assert vu.poster_frame_index(8) == 4 + assert vu.poster_frame_index(16) == 8 + assert vu.poster_frame_index(1) == 0 + + def test_slices_a_per_frame_mask_stack(self): + masks = np.arange(8 * 4 * 4).reshape(8, 4, 4) + got = vu.select_frame_annotation(masks, 8) + assert got.shape == (4, 4) + assert (got == masks[4]).all() + + def test_leaves_a_single_2d_mask_alone(self): + mask = np.zeros((4, 4)) + assert vu.select_frame_annotation(mask, 8).shape == (4, 4) + + def test_does_not_slice_a_colour_mask(self): + """[H, W, C] is 3-D but NOT per-frame — slicing it would take a row.""" + colour = np.zeros((4, 4, 3)) + assert vu.select_frame_annotation(colour, 8).shape == (4, 4, 3) + + def test_slices_per_frame_box_arrays(self): + boxes = np.stack([np.full((1, 6), t) for t in range(8)]) # [T, N, 6] + got = vu.select_frame_annotation(boxes, 8) + assert got.shape == (1, 6) + assert int(got[0, 0]) == 4 + + def test_slices_a_per_frame_list(self): + boxes = [np.full((2, 6), t) for t in range(8)] + assert int(vu.select_frame_annotation(boxes, 8)[0, 0]) == 4 + + def test_image_datasets_are_untouched(self): + """T == 0 is how an image dataset reports itself; nothing may change.""" + arr = np.arange(8 * 4 * 4).reshape(8, 4, 4) + assert vu.select_frame_annotation(arr, 0).shape == arr.shape + assert vu.select_frame_annotation(arr, 1).shape == arr.shape + + def test_none_and_mismatched_lengths_pass_through(self): + assert vu.select_frame_annotation(None, 8) is None + # A list that is not length T is not per-frame. + other = [1, 2, 3] + assert vu.select_frame_annotation(other, 8) is other + + +class TestHasPlayableMedia: + """media_kind must be enough; task_type is not the only way to be video.""" + + class _VideoCls: + media_kind = "video" + task_type = "classification" + + class _Image: + task_type = "classification" + + def test_media_kind_makes_a_classifier_playable(self): + assert vu.has_playable_media(self._VideoCls(), "classification") + + def test_generation_task_is_playable_without_media_kind(self): + assert vu.has_playable_media(self._Image(), "video_generation") + assert vu.has_playable_media(self._Image(), "audio_generation") + + def test_plain_image_dataset_is_not_playable(self): + assert not vu.has_playable_media(self._Image(), "classification") + assert not vu.has_playable_media(self._Image(), "segmentation") + + +class TestPosterFrameHook: + """A dataset may supply the grid's still without decoding a clip. + + Measured on a 40 s 720p corpus: 283 ms to decode a 16-frame window vs + 158 ms for one frame, so roughly half the grid's image time went on + frames nobody looks at. + """ + + class _WithHook: + media_kind = "video" + num_frames = 8 + + def get_poster_frame(self, index): + return np.full((4, 6, 3), int(index), dtype=np.uint8) + + class _Declining: + media_kind = "video" + + def get_poster_frame(self, index): + return None # cannot cheaply seek — caller falls back + + class _Broken: + media_kind = "video" + + def get_poster_frame(self, index): + raise RuntimeError("boom") + + class _WrongShape: + media_kind = "video" + + def get_poster_frame(self, index): + return np.zeros((8, 4, 6, 3), dtype=np.uint8) # a whole clip + + class _NoHook: + media_kind = "video" + + def test_returns_the_frame_the_dataset_supplies(self): + got = vu.load_poster_frame_direct(self._WithHook(), 3) + assert got is not None and got.shape == (4, 6, 3) + assert int(got[0, 0, 0]) == 3 + + def test_dataset_without_the_hook_falls_back(self): + assert vu.load_poster_frame_direct(self._NoHook(), 0) is None + + def test_declining_falls_back(self): + assert vu.load_poster_frame_direct(self._Declining(), 0) is None + + def test_a_raising_hook_falls_back_instead_of_propagating(self): + assert vu.load_poster_frame_direct(self._Broken(), 0) is None + + def test_a_wrong_shaped_poster_is_refused(self): + # Drawn, not raised, if accepted — so refuse rather than guess. + assert vu.load_poster_frame_direct(self._WrongShape(), 0) is None diff --git a/tests/gRPC/test_get_media.py b/tests/gRPC/test_get_media.py index 8909e6dc..5f4197a3 100644 --- a/tests/gRPC/test_get_media.py +++ b/tests/gRPC/test_get_media.py @@ -69,6 +69,7 @@ class _StubService: """Minimal stand-in exposing only what GetMedia touches.""" _MEDIA_CHUNK_BYTES = DataService._MEDIA_CHUNK_BYTES _MEDIA_CACHE_ENTRIES = DataService._MEDIA_CACHE_ENTRIES + _MEDIA_CACHE_BYTES = DataService._MEDIA_CACHE_BYTES GetMedia = DataService.GetMedia _stream_media = DataService._stream_media _media_cache_get = DataService._media_cache_get @@ -432,3 +433,45 @@ def test_chunk_sizes_respect_the_configured_bound(): assert all(len(c.data) <= 2048 for c in chunks) assert sum(len(c.data) for c in chunks) == chunks[0].total_bytes + + +# --------------------------------------------------------------------------- +# Cache budget — entries AND bytes +# --------------------------------------------------------------------------- +def test_media_cache_evicts_on_byte_budget_before_entry_count(): + """A few large clips must not pin memory just because the count fits. + + Clip sizes vary by orders of magnitude between datasets, so the entry + count alone cannot bound this cache: 32 entries of 40 MB each would be + 1.2 GB held inside the training process. + """ + stub = _StubService(_FakeVideoDataset()) + stub._MEDIA_CACHE_ENTRIES = 100 # deliberately not the binding limit + stub._MEDIA_CACHE_BYTES = 10 * 1024 # 10 KB + + for i in range(8): + stub._media_cache_put(i, {"data": b"x" * 4096}) # 4 KB each + + total = sum(len(v["data"]) for v in stub._media_cache.values()) + assert total <= stub._MEDIA_CACHE_BYTES, total + assert len(stub._media_cache) < 8 # byte budget evicted, count did not + + +def test_media_cache_keeps_the_entry_just_inserted(): + """Whoever just put a clip in is about to stream it.""" + stub = _StubService(_FakeVideoDataset()) + stub._MEDIA_CACHE_ENTRIES = 1 + stub._MEDIA_CACHE_BYTES = 1 # smaller than any real payload + + stub._media_cache_put("only", {"data": b"x" * 4096}) + + assert "only" in stub._media_cache + + +def test_media_cache_holds_the_studio_prefetch_window(): + """The default must cover the modal's 6-neighbour window. + + A cache smaller than the window evicts clips the studio has just + prefetched, so every step through a video dataset re-muxes on the backend. + """ + assert DataService._MEDIA_CACHE_ENTRIES >= 7 diff --git a/tests/general/test_example_fashion_mnist_umap.py b/tests/general/test_example_fashion_mnist_umap.py new file mode 100644 index 00000000..e4dea6ac --- /dev/null +++ b/tests/general/test_example_fashion_mnist_umap.py @@ -0,0 +1,499 @@ +"""The Fashion-MNIST "live UMAP, then a PCA of your own" example, and the docs +that point at it. + +``examples/PyTorch/wl-fashion-mnist-umap/main.py`` drives a run in stages: +train with the built-in parametric UMAP, keep a checkpoint, reload it, discard +the lowest-loss training samples, project what is left with a PCA written in +plain torch, compare the two layouts, retrain. + +Three things are held here: + +* the building blocks (loss cut, PCA, neighbour purity, layout comparison) on + small synthetic inputs; +* the whole script, run for real on a synthetic Fashion-MNIST (no download), + against the real ledger, checkpoint manager and projection; +* the docs: the page and the gallery entry. +""" + +import contextlib +import importlib.util +import io +import os +import re +import shutil +import tempfile +import unittest +from pathlib import Path +from unittest import mock + +import numpy as np +import pandas as pd +import torch +from PIL import Image + +import weightslab as wl +from weightslab import src as wl_src +from weightslab.backend import ledgers +from weightslab.projection import detach_projection +from weightslab.projection.registry import clear_registry, known_prefixes + +REPO_ROOT = Path(__file__).resolve().parents[2] +EXAMPLE_DIR = REPO_ROOT / "weightslab" / "examples" / "PyTorch" / "wl-fashion-mnist-umap" +DOCS = REPO_ROOT / "docs" + +LOSS = "signals//train-loss-CE" + + +def load_example(): + """Import the example's main.py by path (its directory is not a package).""" + spec = importlib.util.spec_from_file_location("wl_fashion_mnist_umap_example", + EXAMPLE_DIR / "main.py") + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +example = load_example() + + +# ============================================================================= +# The building blocks +# ============================================================================= +def sample_rows(n=10, **columns): + """A per-sample table the way ``sample_table()`` returns it.""" + return pd.DataFrame(columns, index=pd.Index(range(n), name="sample_id")) + + +class TestLowestLossIds(unittest.TestCase): + + def test_picks_the_lowest_fraction(self): + table = sample_rows(10, **{LOSS: [9, 1, 8, 2, 7, 3, 6, 4, 5, 0.5]}) + self.assertEqual(sorted(example.lowest_loss_ids(table, 0.3)), [1, 3, 9]) + + def test_returns_plain_ints(self): + table = sample_rows(4, **{LOSS: [0.4, 0.1, 0.3, 0.2]}) + ids = example.lowest_loss_ids(table, 0.5) + self.assertTrue(all(type(i) is int for i in ids)) + self.assertEqual(ids, [1, 3]) # lowest first + + def test_a_sample_with_no_loss_is_never_picked(self): + """Unseen is not easy: it must not be discarded as if the model knew it.""" + table = sample_rows(4, **{LOSS: [np.nan, 0.2, np.nan, 0.1]}) + self.assertEqual(example.lowest_loss_ids(table, 1.0), [3, 1]) + + def test_a_fraction_of_zero_or_a_missing_column_picks_nothing(self): + table = sample_rows(4, **{LOSS: [0.4, 0.1, 0.3, 0.2]}) + self.assertEqual(example.lowest_loss_ids(table, 0), []) + self.assertEqual(example.lowest_loss_ids(table, 0.5, column="signals//nope"), []) + + def test_a_fraction_above_one_means_all_of_them(self): + table = sample_rows(4, **{LOSS: [0.4, 0.1, 0.3, 0.2]}) + self.assertEqual(len(example.lowest_loss_ids(table, 7.0)), 4) + + def test_losses_stored_as_objects_still_compare_as_numbers(self): + table = sample_rows(3, **{LOSS: np.array([10.0, 9.0, 100.0], dtype=object)}) + self.assertEqual(example.lowest_loss_ids(table, 0.34), [1]) # not '100.0' < '9.0' + + +class TestCheckpointFailuresAreLoud(unittest.TestCase): + """A checkpoint that was not written, or not reloaded, must stop the run + there -- not surface two stages later as a confusing result.""" + + def test_a_checkpoint_that_could_not_be_written_is_an_error(self): + manager = mock.MagicMock() + manager.save_model_checkpoint.return_value = None # what the manager gives + with mock.patch.object(example.ledgers, "get_checkpoint_manager", return_value=manager): + with self.assertRaisesRegex(RuntimeError, "could not write a checkpoint"): + example.save_checkpoint() + + def test_a_checkpoint_that_could_not_be_reloaded_is_an_error(self): + manager = mock.MagicMock() + manager.load_state.return_value = False + with mock.patch.object(example.ledgers, "get_checkpoint_manager", return_value=manager): + with self.assertRaisesRegex(RuntimeError, "step 40"): + example.reload_checkpoint(40) + self.assertEqual(manager.load_state.call_args.kwargs["target_step"], 40) + + +class TestPcaCoordinates(unittest.TestCase): + + def setUp(self): + rng = np.random.default_rng(0) + # 400 points in 5-D: almost all variance along two known directions. + scale = np.array([10.0, 4.0, 0.05, 0.05, 0.05]) + self.features = rng.normal(size=(400, 5)) * scale + 100.0 # off-centre on purpose + + def test_shape_for_two_and_three_components(self): + for dim in (2, 3): + coords, _ = example.pca_coordinates(self.features, dim) + self.assertEqual(coords.shape, (400, dim)) + + def test_the_coordinates_are_centred(self): + coords, _ = example.pca_coordinates(self.features, 2) + np.testing.assert_allclose(coords.mean(0), 0.0, atol=1e-3) + + def test_the_variance_share_matches_the_covariance_spectrum(self): + _, share = example.pca_coordinates(self.features, 2) + eigenvalues = np.sort(np.linalg.eigvalsh(np.cov(self.features.T)))[::-1] + np.testing.assert_allclose(share, eigenvalues[:2].sum() / eigenvalues.sum(), rtol=1e-3) + self.assertGreater(share, 0.99) + + def test_the_first_axis_carries_the_most_variance(self): + coords, _ = example.pca_coordinates(self.features, 3) + variances = coords.var(0) + self.assertTrue(variances[0] > variances[1] > variances[2]) + + def test_distances_never_grow(self): + """A projection can only lose information: pairwise distances shrink.""" + coords, _ = example.pca_coordinates(self.features, 2) + original = np.linalg.norm(self.features[0] - self.features[1]) + self.assertLessEqual(np.linalg.norm(coords[0] - coords[1]), original + 1e-4) + + def test_accepts_a_tensor(self): + coords, _ = example.pca_coordinates(torch.as_tensor(self.features), 2) + self.assertEqual(coords.shape, (400, 2)) + + +class TestKnnPurity(unittest.TestCase): + + def test_separated_classes_are_perfectly_pure(self): + coords = np.concatenate([np.zeros((30, 2)), np.full((30, 2), 100.0)]) + coords += np.random.default_rng(0).normal(size=coords.shape) + labels = np.array([0] * 30 + [1] * 30) + self.assertEqual(example.knn_purity(coords, labels, k=5), 1.0) + + def test_unrelated_labels_sit_near_chance(self): + rng = np.random.default_rng(0) + coords = rng.normal(size=(1500, 2)) + labels = rng.integers(0, 10, size=1500) + self.assertLess(example.knn_purity(coords, labels, k=10), 0.2) + + def test_a_point_is_not_its_own_neighbour(self): + """Two points, two classes: if a point counted itself, purity would be 1.""" + coords = np.array([[0.0, 0.0], [1.0, 0.0]]) + self.assertEqual(example.knn_purity(coords, np.array([0, 1]), k=1), 0.0) + + def test_k_larger_than_the_set_is_clipped(self): + coords = np.array([[0.0], [1.0], [2.0]]) + value = example.knn_purity(coords, np.array([0, 0, 0]), k=50) + self.assertEqual(value, 1.0) + + def test_a_single_point_has_no_neighbours(self): + self.assertTrue(np.isnan(example.knn_purity(np.zeros((1, 2)), np.array([0]), k=3))) + + def test_works_across_the_row_blocks_it_computes_in(self): + """More rows than one block (2048): the blocks must line up.""" + rng = np.random.default_rng(1) + n = 2100 + labels = rng.integers(0, 3, size=n) + coords = np.stack([labels * 50.0, np.zeros(n)], axis=1) + rng.normal(size=(n, 2)) + self.assertGreater(example.knn_purity(coords, labels, k=5), 0.99) + + def test_three_dimensional_layouts_work_too(self): + coords = np.concatenate([np.zeros((20, 3)), np.full((20, 3), 50.0)]) + coords += np.random.default_rng(0).normal(size=coords.shape) + self.assertEqual(example.knn_purity(coords, np.array([0] * 20 + [1] * 20), k=4), 1.0) + + +class TestLayoutAndCompare(unittest.TestCase): + + def setUp(self): + n = 12 + rng = np.random.default_rng(0) + a = rng.normal(size=(n, 2)) + b = rng.normal(size=(n, 3)) + a[10:] = np.nan # "a" never placed 10 and 11 + b[:3] = np.nan # "b" never placed 0..2 + discarded = np.zeros(n, dtype=bool) + discarded[5] = True + self.table = sample_rows( + n, target=[0, 1] * 6, discarded=discarded, + **{"signals//a_x": a[:, 0], "signals//a_y": a[:, 1], + "signals//b_x": b[:, 0], "signals//b_y": b[:, 1], "signals//b_z": b[:, 2]}) + + def test_layout_leaves_out_discarded_and_unplaced_samples(self): + ids, coords = example.layout(self.table, "a") + self.assertEqual(ids, [0, 1, 2, 3, 4, 6, 7, 8, 9]) # not 5 (discarded), 10, 11 + self.assertEqual(coords.shape, (9, 2)) + + def test_layout_follows_the_columns_that_exist(self): + ids, coords = example.layout(self.table, "b") + self.assertEqual(coords.shape[1], 3) + self.assertNotIn(5, ids) + self.assertEqual(ids[0], 3) + + def test_an_unknown_prefix_places_nothing(self): + ids, coords = example.layout(self.table, "nope") + self.assertEqual(ids, []) + self.assertEqual(coords.size, 0) + + def test_compare_scores_the_samples_every_layout_places(self): + scores, shared = example.compare(self.table, ["a", "b"], k=2) + self.assertEqual(shared, 6) # 3,4,6,7,8,9 + self.assertEqual(set(scores), {"a", "b"}) + for value in scores.values(): + self.assertTrue(0.0 <= value <= 1.0) + + +# ============================================================================= +# The whole script, for real +# ============================================================================= +class FakeFashionMNIST: + """``torchvision.datasets.FashionMNIST``'s interface -- ``(PIL image, label)`` + items and ``.targets`` -- with no download. Class ``c`` is a bright band at + rows ``2c+3 .. 2c+5``, so even a few steps tell the classes apart.""" + + def __init__(self, root=None, train=True, download=False, transform=None): + count = 1200 if train else 300 + rng = np.random.default_rng(0 if train else 1) + self.targets = torch.as_tensor(np.arange(count) % 10) + self.images = rng.integers(0, 40, size=(count, 28, 28), dtype=np.uint8) + for index, label in enumerate(self.targets.tolist()): + self.images[index, 2 * label + 3:2 * label + 6, :] += 150 + + def __len__(self): + return len(self.targets) + + def __getitem__(self, index): + return Image.fromarray(self.images[index]), int(self.targets[index]) + + +TRAIN_SIZE, TEST_SIZE = 400, 100 +STEPS_INITIAL, STEPS_RETRAIN, DISCARD = 24, 12, 0.25 + + +def tiny_parameters(root): + return { + "experiment_name": "test_fashion_mnist_umap", + "device": "cpu", + "root_log_dir": root, + "data_root": os.path.join(root, "data"), + "projection": {"every_n_steps": 4, "graph_size": 128}, + "exploration": {"steps_initial": STEPS_INITIAL, "steps_retrain": STEPS_RETRAIN, + "discard_fraction": DISCARD, "knn_k": 5}, + "eval_full_to_train_steps_ratio": 12, + "experiment_dump_to_train_steps_ratio": 10 ** 9, # only the explicit checkpoints + "tqdm_display": False, + "start_timeout": 0, + "serving_grpc": False, + "ledger_enable_h5_persistence": False, + "data": {"train_loader": {"batch_size": 32, "max_samples": TRAIN_SIZE}, + "test_loader": {"batch_size": 50, "max_samples": TEST_SIZE}}, + } + + +class TestTheStagedRun(unittest.TestCase): + """One real run of ``run()``; each test reads a different part of it.""" + + @classmethod + def setUpClass(cls): + ledgers.clear_all() + detach_projection() + clear_registry() + cls.root = tempfile.mkdtemp(prefix="wl-fmnist-umap-") + cls.saved_checkpoints = {} + real_save = example.save_checkpoint + + def spy(): + step = real_save() + # What the train table looked like when this checkpoint was written. + cls.saved_checkpoints.setdefault( + step, int(example.sample_table()["nb_seen"].astype(float).sum())) + return step + + real_cut = example.lowest_loss_ids + + def cut_spy(table, fraction, **kwargs): + ids = real_cut(table, fraction, **kwargs) + cls.cut = (table.copy(), ids) # the table the cut was made from + return ids + + cls.served = mock.MagicMock() + try: + with mock.patch.object(example.datasets, "FashionMNIST", FakeFashionMNIST), \ + mock.patch.object(example, "save_checkpoint", spy), mock.patch.object(example, "lowest_loss_ids", cut_spy), \ + mock.patch.object(wl, "serve", cls.served), \ + contextlib.redirect_stdout(io.StringIO()) as out: + cls.summary = example.run(tiny_parameters(cls.root)) + cls.output = out.getvalue() + # The end of the run, kept before the reload below rewinds it. + cls.table = example.sample_table() + cls.prefixes = set(known_prefixes()) + cls.reload = cls.reload_the_checkpoint() + except BaseException: + cls.tearDownClass() + raise + + @classmethod + def reload_the_checkpoint(cls): + """Train has gone on to step 36; reload step 24 and read what came back.""" + model = ledgers.get_model() + found = {"age_before": model.get_age()} + + # The UMAP encoder and the custom projections are saved beside each + # checkpoint under the same step number: the weights are the one file + # that holds a model_state_dict. + manager = ledgers.get_checkpoint_manager() + files = sorted(Path(manager.models_dir).rglob(f"*_step_{STEPS_INITIAL:06d}.pt")) + loaded = [torch.load(f, map_location="cpu", weights_only=False) for f in files] + weights = [c["model_state_dict"] for c in loaded + if isinstance(c, dict) and "model_state_dict" in c] + found["checkpointed_weights"] = weights[-1] if weights else None + + found["age_after"] = example.reload_checkpoint(STEPS_INITIAL) + found["weights_after"] = {k: v.detach().cpu().clone() + for k, v in ledgers.get_model().state_dict().items() + if torch.is_tensor(v)} + found["seen_after"] = int(example.sample_table()["nb_seen"].astype(float).sum()) + return found + + @classmethod + def tearDownClass(cls): + wl.pause_training() # a fresh process starts paused + detach_projection() + clear_registry() + ledgers.clear_all() + wl_src.DATAFRAME_M = None # clear_all orphans the cached handle + shutil.rmtree(cls.root, ignore_errors=True) + + # -- the stages ------------------------------------------------------------ + def test_trains_to_the_checkpoint_step_and_keeps_a_checkpoint(self): + self.assertEqual(self.summary["checkpoint_step"], STEPS_INITIAL) + + def test_the_reload_lands_on_the_checkpoint(self): + self.assertEqual(self.summary["reloaded_age"], STEPS_INITIAL) + + def test_discards_the_requested_share_of_the_training_samples(self): + self.assertEqual(self.summary["discarded"], int(TRAIN_SIZE * DISCARD)) + table = self.table + self.assertEqual(int(table["discarded"].astype(bool).sum()), self.summary["discarded"]) + self.assertEqual(len(table), TRAIN_SIZE) # train split only + + def test_what_was_discarded_is_what_the_model_knew_best(self): + """At the moment of the cut, every discarded sample has a loss no larger + than every sample that stayed. (Not checkable on the end state: kept + samples keep training afterwards, discarded ones are frozen.)""" + table, ids = self.cut + loss = pd.to_numeric(table[LOSS], errors="coerce") + gone, kept = loss.loc[ids], loss.drop(index=ids).dropna() + self.assertEqual(len(ids), self.summary["discarded"]) + self.assertLessEqual(gone.max(), kept.min()) + + def test_the_pca_covers_exactly_what_was_kept(self): + """The loader skips discarded samples, and the sweep sees the WHOLE split + even though training was mid-epoch on it (the stateful-loader trap).""" + kept = TRAIN_SIZE - self.summary["discarded"] + self.assertEqual(self.summary["pca_samples"], kept) + table = self.table + placed = table[f"signals//pca2d_step{STEPS_INITIAL}_x"].notna() + self.assertEqual(int(placed.sum()), kept) + self.assertFalse((placed & table["discarded"].astype(bool)).any()) + + def test_the_pca_is_two_dimensional(self): + table = self.table + self.assertIn(f"signals//pca2d_step{STEPS_INITIAL}_y", table.columns) + self.assertNotIn(f"signals//pca2d_step{STEPS_INITIAL}_z", table.columns) + + def test_the_live_umap_placed_the_training_samples(self): + table = self.table + for axis in "xyz": + self.assertGreater(int(table[f"signals//umap_{axis}"].notna().sum()), TRAIN_SIZE // 2) + + def test_both_layouts_are_scored(self): + for value in (*self.summary["purity"].values(), self.summary["purity_features"]): + self.assertTrue(0.0 <= value <= 1.0, value) + self.assertEqual(set(self.summary["purity"]), {"umap", f"pca2d_step{STEPS_INITIAL}"}) + + def test_the_figure_is_written(self): + figure = Path(self.summary["log_dir"]) / "umap_vs_pca.png" + self.assertTrue(figure.is_file()) + self.assertGreater(figure.stat().st_size, 5_000) + + def test_the_picker_lists_every_projection(self): + self.assertIn("umap", self.prefixes) + self.assertIn(f"pca2d_step{STEPS_INITIAL}", self.prefixes) + self.assertIn(f"pca2d_step{STEPS_INITIAL + STEPS_RETRAIN}", self.prefixes) + + # -- the retraining ---------------------------------------------------------- + def test_neither_sweeps_nor_projections_age_the_model(self): + """Five sweeps ran (PCA, PCA again, the UMAP's own placements): the model + trained exactly 24 + 12 steps, not one more.""" + self.assertEqual(self.summary["final_age"], STEPS_INITIAL + STEPS_RETRAIN) + + def test_the_second_pca_is_a_new_projection_not_an_overwrite(self): + after = f"pca2d_step{STEPS_INITIAL + STEPS_RETRAIN}" + self.assertIn(after, self.summary["purity_after"]) + table = self.table + self.assertEqual(int(table[f"signals//{after}_x"].notna().sum()), + TRAIN_SIZE - self.summary["discarded"]) + self.assertEqual(int(table[f"signals//pca2d_step{STEPS_INITIAL}_x"].notna().sum()), + TRAIN_SIZE - self.summary["discarded"]) + + # -- the reload ------------------------------------------------------------- + def test_reload_brings_back_the_age_of_the_checkpoint(self): + self.assertEqual(self.reload["age_before"], STEPS_INITIAL + STEPS_RETRAIN) + self.assertEqual(self.reload["age_after"], STEPS_INITIAL) + + def test_reload_brings_back_the_weights_in_the_checkpoint_file(self): + saved = self.reload["checkpointed_weights"] + self.assertIsNotNone(saved, "no weights checkpoint was written at the checkpoint step") + checked = 0 + for name, value in saved.items(): + if torch.is_tensor(value) and value.dtype.is_floating_point: + self.assertTrue(torch.equal(value, self.reload["weights_after"][name]), name) + checked += 1 + self.assertGreater(checked, 4) + + def test_reload_rewinds_the_sample_counters_to_the_checkpoint(self): + """The per-sample seen-counts go back too, instead of describing steps + the reload just undid. They land within ONE batch of the checkpoint, not + exactly on it: signals are logged at model age - 1 (see the TODO on + ``src._get_step``), so a restore to step N still includes the signals the + (N+1)-th batch logged at step N.""" + at_checkpoint = self.saved_checkpoints[STEPS_INITIAL] + at_end = int(self.table["nb_seen"].astype(float).sum()) + after = self.reload["seen_after"] + self.assertLess(after, at_end) + self.assertGreaterEqual(after, at_checkpoint) + self.assertLessEqual(after - at_checkpoint, 32) # one batch + + # -- wiring ------------------------------------------------------------------ + def test_serves_the_studio_backend_as_configured(self): + self.served.assert_called_once_with(serving_grpc=False) + + def test_announces_each_stage(self): + for stage in ("[1. Train", "[2. Reload", "[3. Discard", "[4. PCA", "[5. Compare", + "[6. Retrain"): + self.assertIn(stage, self.output) + + +# ============================================================================= +# The docs +# ============================================================================= +@unittest.skipUnless(DOCS.is_dir(), "the docs are not part of this checkout") +class TestTheDocs(unittest.TestCase): + + def read(self, *parts): + return (DOCS.joinpath(*parts)).read_text(encoding="utf-8") + + def test_the_example_has_a_page_in_the_pytorch_section(self): + page = self.read("examples", "pytorch", "fashion_mnist_umap.rst") + self.assertIn("weightslab/examples/PyTorch/wl-fashion-mnist-umap/main.py", page) + self.assertRegex(self.read("examples", "pytorch", "index.rst"), + r"(?m)^\s+fashion_mnist_umap\s*$") + + def test_the_usecase_page_points_at_the_script(self): + page = self.read("examples", "usecases", "projection_fashion_mnist.rst") + self.assertIn("PyTorch/wl-fashion-mnist-umap", page) + + def test_the_gallery_lists_the_example(self): + gallery = self.read("_static", "examples-gallery.js") + entry = re.search(r"\{[^{}]*url:\s*'examples/pytorch/fashion_mnist_umap\.html'[^{}]*\}", + gallery) + self.assertIsNotNone(entry, "no gallery entry for the example") + self.assertIn("badge: 'PyTorch'", entry.group(0)) + self.assertTrue((DOCS / "examples" / "pytorch" / "fashion_mnist_umap.rst").is_file()) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/general/test_parametric_umap.py b/tests/general/test_parametric_umap.py new file mode 100644 index 00000000..fc99b979 --- /dev/null +++ b/tests/general/test_parametric_umap.py @@ -0,0 +1,1574 @@ +"""Parametric UMAP: the projection itself, and the view-box serving around it.""" + +import os +import time +import unittest + +import numpy as np +import pandas as pd +import torch as th +import torch.nn as nn + +from weightslab import src as wl_src +from weightslab.projection import ( + ENV_ENABLED, + ProjectionTracker, + attach_projection, + detach_projection, + find_ab_params, + first_tensor, + flatten_features, + get_tracker, + membership_high_dim, + pick_embedding_layer, + projection_enabled, +) +from weightslab.proto import experiment_service_pb2 as pb2 +from weightslab.trainer.services import projection_service as ps + + +class TinyNet(nn.Module): + """conv-ish trunk -> features -> classifier, so the auto-picked layer has + an obvious right answer (the classifier's INPUT, i.e. `feat`'s output).""" + + def __init__(self, in_dim=16, feat_dim=8, classes=3): + super().__init__() + self.feat = nn.Linear(in_dim, feat_dim) + self.act = nn.ReLU() + self.classifier = nn.Linear(feat_dim, classes) + + def forward(self, x): + return self.classifier(self.act(self.feat(x))) + + +def clustered(n_per=40, dim=16, clusters=3, seed=0): + """Well-separated blobs -- structure the projection is expected to keep.""" + rng = np.random.default_rng(seed) + centers = rng.normal(scale=6.0, size=(clusters, dim)) + xs, ys = [], [] + for c in range(clusters): + xs.append(centers[c] + rng.normal(scale=0.35, size=(n_per, dim))) + ys.append(np.full(n_per, c)) + return (th.tensor(np.concatenate(xs), dtype=th.float32), + np.concatenate(ys)) + + +class TestAbParams(unittest.TestCase): + def test_matches_umap_learn_reference(self): + a, b = find_ab_params(spread=1.0, min_dist=0.1) + # umap-learn's scipy fit for these inputs. + self.assertAlmostEqual(a, 1.5769434, delta=0.08) + self.assertAlmostEqual(b, 0.8950608, delta=0.05) + + def test_smaller_min_dist_tightens_the_curve(self): + a_tight, _ = find_ab_params(spread=1.0, min_dist=0.0) + a_loose, _ = find_ab_params(spread=1.0, min_dist=0.5) + self.assertGreater(a_tight, a_loose) + + +class TestMembership(unittest.TestCase): + def test_rows_sum_to_log2_k_before_symmetrisation(self): + feats, _ = clustered(n_per=20) + k = 10 + # Re-derive the one-sided memberships the way the function does, since + # the symmetrised output deliberately no longer sums to log2(k). + dist = th.cdist(feats, feats) + dist = dist + th.eye(len(feats)) * 1e12 + knn_d, _ = th.topk(dist, k, dim=1, largest=False) + mu = membership_high_dim(feats, k) + self.assertEqual(mu.shape, (len(feats), len(feats))) + # Symmetric, in [0,1], zero on the diagonal. + self.assertTrue(th.allclose(mu, mu.t(), atol=1e-6)) + self.assertGreaterEqual(float(mu.min()), 0.0) + self.assertLessEqual(float(mu.max()), 1.0 + 1e-6) + self.assertLess(float(mu.diagonal().abs().max()), 1e-6) + + def test_near_neighbours_score_above_far_ones(self): + feats, labels = clustered(n_per=25, clusters=3) + mu = membership_high_dim(feats, 15).numpy() + same = np.equal.outer(labels, labels) + np.fill_diagonal(same, False) + self.assertGreater(mu[same].mean(), mu[~same].mean()) + + +class TestTracker(unittest.TestCase): + def test_fitting_preserves_cluster_structure(self): + """The point of the whole feature: separated in feature space should + come out separated in 3-D.""" + th.manual_seed(0) + feats, labels = clustered(n_per=40, dim=16, clusters=3) + tracker = ProjectionTracker(out_dim=3, n_neighbors=12, inner_steps=1) + tracker._ensure_encoder(feats.shape[1], feats.device) + target = membership_high_dim(feats, 12) + + first = None + for step in range(220): + tracker._optimizer.zero_grad(set_to_none=True) + loss = tracker._umap_loss(feats, target) + loss.backward() + tracker._optimizer.step() + if step == 0: + first = loss.item() + self.assertLess(loss.item(), first) + + with th.no_grad(): + emb = tracker._encoder(feats).numpy() + within, between = [], [] + for i in range(3): + centroid = emb[labels == i].mean(axis=0) + within.append(np.linalg.norm(emb[labels == i] - centroid, axis=1).mean()) + for j in range(i + 1, 3): + between.append(np.linalg.norm(centroid - emb[labels == j].mean(axis=0))) + self.assertGreater(np.mean(between), 2.0 * np.mean(within)) + + def test_hook_captures_penultimate_features_not_logits(self): + model = TinyNet(in_dim=16, feat_dim=8, classes=3) + layer, use_input, feature_last = pick_embedding_layer(model) + self.assertIs(layer, model.classifier) + self.assertTrue(use_input) + self.assertTrue(feature_last) + + tracker = ProjectionTracker() + tracker.attach(layer, use_input=use_input, feature_last=feature_last) + try: + model(th.randn(5, 16)) + self.assertIsNotNone(tracker._pending) + # 8 (features), not 3 (classes). + self.assertEqual(tuple(tracker._pending.shape), (5, 8)) + finally: + tracker.detach() + + def test_hook_never_backpropagates_into_the_model(self): + """Instrumentation must not touch the training graph.""" + model = TinyNet() + tracker = ProjectionTracker() + tracker.attach(model.classifier, use_input=True) + try: + model(th.randn(6, 16)) + self.assertFalse(tracker._pending.requires_grad) + finally: + tracker.detach() + + def test_conv_shaped_features_are_pooled_to_2d(self): + conv = nn.Conv2d(3, 7, 3, padding=1) + tracker = ProjectionTracker() + tracker.attach(conv, use_input=False) + try: + conv(th.randn(4, 3, 8, 8)) + self.assertEqual(tuple(tracker._pending.shape), (4, 7)) + finally: + tracker.detach() + + def test_mismatched_batch_ids_are_refused(self): + """Pairing is positional; a length mismatch would silently attach one + sample's features to another's id, so it must be dropped.""" + model = TinyNet() + tracker = ProjectionTracker(every_n_steps=1) + tracker.attach(model.classifier) + try: + model(th.randn(8, 16)) + self.assertFalse(tracker.observe_batch(["a", "b"], step=1)) + finally: + tracker.detach() + + def test_a_layer_that_never_sees_one_row_per_sample_says_so(self): + """A per-point layer mismatches on EVERY batch; dropping each one + without a word left the board empty with no explanation.""" + model = TinyNet() + tracker = ProjectionTracker(every_n_steps=1) + tracker.attach(model.classifier) + tracker.layer_name = "classifier" + try: + with self.assertLogs("weightslab.projection.parametric_umap", level="WARNING") as logs: + for step in range(1, 4): + model(th.randn(8, 16)) + self.assertFalse(tracker.observe_batch(["a", "b"], step=step)) + finally: + tracker.detach() + self.assertEqual(tracker.mismatched_batches, 3) + self.assertEqual(tracker.last_mismatch, (8, 2)) + self.assertEqual(tracker.stats()["mismatched_batches"], 3) + self.assertEqual(len(logs.output), 1) + self.assertIn("8 feature rows for a batch of 2 samples", logs.output[0]) + self.assertIn("'classifier'", logs.output[0]) + + +class TestSaveSignalsCoercion(unittest.TestCase): + """save_signals must never create a column it then writes nothing into.""" + + class _Stub: + def __init__(self): + self.calls = [] + + def enqueue_batch(self, **kwargs): + self.calls.append(kwargs) + + def setUp(self): + self._saved = wl_src.DATAFRAME_M + self.stub = self._Stub() + wl_src.DATAFRAME_M = self.stub + + def tearDown(self): + wl_src.DATAFRAME_M = self._saved + + def test_numpy_arrays_reach_the_dataframe(self): + """The regression that motivated this: a bare ndarray used to normalize + to None, so the column appeared and stayed empty forever.""" + values = np.array([1.0, 2.0, 3.0], dtype=np.float32) + wl_src.save_signals(signals={"metric": values}, batch_ids=[1, 2, 3], step=0) + losses = self.stub.calls[0]["losses"] + self.assertIn("signals//metric", losses) + self.assertIsNotNone(losses["signals//metric"]) + np.testing.assert_allclose(losses["signals//metric"], values) + + def test_tensors_still_reach_the_dataframe(self): + wl_src.save_signals(signals={"metric": th.tensor([1.0, 2.0])}, + batch_ids=[1, 2], step=0) + np.testing.assert_allclose( + self.stub.calls[0]["losses"]["signals//metric"], [1.0, 2.0]) + + def test_lists_still_reach_the_dataframe(self): + wl_src.save_signals(signals={"metric": [1.0, 2.0]}, batch_ids=[1, 2], step=0) + self.assertIsNotNone(self.stub.calls[0]["losses"]["signals//metric"]) + + def test_multidim_values_reduce_to_one_per_sample(self): + wl_src.save_signals(signals={"metric": np.ones((3, 4), dtype=np.float32)}, + batch_ids=[1, 2, 3], step=0) + self.assertEqual(self.stub.calls[0]["losses"]["signals//metric"].shape, (3,)) + + def test_unsupported_type_raises_instead_of_writing_nothing(self): + with self.assertRaises(TypeError) as ctx: + wl_src.save_signals(signals={"metric": {"not": "a signal"}}, + batch_ids=[1, 2], step=0) + message = str(ctx.exception) + self.assertIn("metric", message) # names the offending signal + self.assertIn("batch_ids", message) # says what a good value is + self.assertEqual(self.stub.calls, []) # nothing half-written + + +class TestAnyTorchModel(unittest.TestCase): + """The hook has to land on a usable (B, F) embedding for whatever the user + actually wrapped -- not just for a toy classifier.""" + + def test_flatten_keeps_the_feature_axis_for_each_layout(self): + # Linear INPUT: features are the last axis whatever the rank, because + # that is the axis Linear acts on. + self.assertEqual( + tuple(flatten_features(th.randn(4, 7, 16), feature_last=True).shape), (4, 16)) + self.assertEqual( + tuple(flatten_features(th.randn(4, 16), feature_last=True).shape), (4, 16)) + # Conv OUTPUT: channels-first, so pool the spatial axes. + self.assertEqual( + tuple(flatten_features(th.randn(4, 16, 8, 8), feature_last=False).shape), (4, 16)) + self.assertEqual( + tuple(flatten_features(th.randn(4, 16, 8, 8, 3), feature_last=False).shape), (4, 16)) + + def test_first_tensor_unwraps_real_model_outputs(self): + tensor = th.randn(2, 3) + self.assertIs(first_tensor(tensor), tensor) + self.assertIs(first_tensor((tensor, None)), tensor) + self.assertIs(first_tensor([None, [tensor]]), tensor) + self.assertIs(first_tensor({"last_hidden_state": tensor}), tensor) + self.assertIs(first_tensor({"anything": tensor}), tensor) + self.assertIsNone(first_tensor({"nope": "text"})) + + def _assert_projects(self, model, sample_input, expected_dim): + layer, use_input, feature_last = pick_embedding_layer(model) + self.assertIsNotNone(layer) + tracker = ProjectionTracker(every_n_steps=1) + tracker.attach(layer, use_input=use_input, feature_last=feature_last) + try: + model(sample_input) + self.assertIsNotNone(tracker._pending, "hook captured nothing") + self.assertEqual(tracker._pending.ndim, 2) + self.assertEqual(tracker._pending.shape[0], sample_input.shape[0]) + self.assertEqual(tracker._pending.shape[1], expected_dim) + finally: + tracker.detach() + + def test_tabular_mlp(self): + model = nn.Sequential(nn.Linear(12, 32), nn.ReLU(), nn.Linear(32, 2)) + self._assert_projects(model, th.randn(8, 12), 32) + + def test_vision_cnn(self): + model = nn.Sequential( + nn.Conv2d(3, 8, 3, padding=1), nn.ReLU(), nn.AdaptiveAvgPool2d(1), + nn.Flatten(), nn.Linear(8, 5)) + self._assert_projects(model, th.randn(6, 3, 16, 16), 8) + + def test_sequence_model_pools_tokens_not_channels(self): + """(B, T, C) into the head: keep C, pool T. Getting this backwards + would hand the projection a per-token scalar.""" + class Seq(nn.Module): + def __init__(self): + super().__init__() + self.encoder = nn.TransformerEncoderLayer( + d_model=32, nhead=4, dim_feedforward=64, batch_first=True) + self.head = nn.Linear(32, 3) + + def forward(self, x): + return self.head(self.encoder(x)) + + self._assert_projects(Seq(), th.randn(5, 9, 32), 32) + + def test_point_cloud_detector_hooks_its_head_not_the_per_point_net(self): + """PointPillars (the 2-D / 3-D LiDAR examples): the only Linear is the + per-point feature net, registered FIRST and run on every point of the + batch at once -- one row per point. The head is a conv on the stacked + BEV canvases, one per sample, and that is what has to be hooked: the + Linear's rows never matched the batch, so every batch was dropped.""" + class Pillars(nn.Module): + def __init__(self): + super().__init__() + self.pfn = nn.Sequential(nn.Linear(4, 8), nn.ReLU()) + self.backbone = nn.Sequential(nn.Conv2d(8, 16, 3, padding=1), nn.ReLU()) + self.head = nn.Conv2d(16, 5, 1) + + def forward(self, points): # (B, N, 4) + b, n, _ = points.shape + per_point = self.pfn(points.reshape(b * n, 4)) + canvas = per_point.reshape(b, n, 8).mean(1)[:, :, None, None] + return self.head(self.backbone(canvas.expand(b, 8, 4, 4).contiguous())) + + model = Pillars() + layer, use_input, feature_last = pick_embedding_layer(model) + self.assertIs(layer, model.head) + self.assertTrue(use_input) + self.assertFalse(feature_last) + self._assert_projects(model, th.randn(3, 50, 4), 16) + + def test_fully_convolutional_segmentation_head(self): + """No Linear anywhere: the last CONV is the head, and its INPUT is the + representation. Reading its output would give 4 class-logit channels -- + well-shaped, and completely the wrong thing to lay out.""" + model = nn.Sequential( + nn.Conv2d(3, 16, 3, padding=1), nn.ReLU(), nn.Conv2d(16, 4, 1)) + layer, use_input, feature_last = pick_embedding_layer(model) + self.assertIsInstance(layer, nn.Conv2d) + self.assertTrue(use_input, "must read the conv head's INPUT, not its logits") + self.assertFalse(feature_last, "conv activations are channels-first") + # 16 (decoder features), not 4 (classes). + self._assert_projects(model, th.randn(4, 3, 12, 12), 16) + + def test_model_returning_a_tuple(self): + class TupleOut(nn.Module): + def __init__(self): + super().__init__() + self.trunk = nn.Linear(8, 16) + self.head = nn.Linear(16, 2) + + def forward(self, x): + hidden = self.trunk(x) + return self.head(hidden), hidden + + self._assert_projects(TupleOut(), th.randn(3, 8), 16) + + def test_eval_batches_are_placed_but_not_fit(self): + """Held-out samples belong in the picture; letting them shape the + projection would leak the eval set into it.""" + model = TinyNet() + tracker = ProjectionTracker(every_n_steps=1) + tracker.attach(model.classifier) + try: + model.train() + model(th.randn(8, 16)) + tracker.observe_batch(list(range(8)), step=1) + fits_after_train = tracker.steps_trained + self.assertGreater(fits_after_train, 0) + + model.eval() + with th.no_grad(): + model(th.randn(8, 16)) + tracker.observe_batch(list(range(8, 16)), step=2) + self.assertEqual(tracker.steps_trained, fits_after_train) + finally: + tracker.detach() + + def test_survives_a_live_architecture_edit(self): + """Editing a running model's width is a headline WeightsLab feature. An + encoder pinned to the old width raises on every batch afterwards, and + the broad except around the fit turns that into a projection that + silently stops updating from the user's first prune onward.""" + model = TinyNet(in_dim=16, feat_dim=8) + tracker = ProjectionTracker(every_n_steps=1) + tracker.attach(model.classifier) + try: + model(th.randn(8, 16)) + tracker.observe_batch(list(range(8)), step=1) + self.assertEqual(tracker.steps_trained, 1) + self.assertEqual(tracker.stats()["feature_dim"], 8) + + # Prune the feature layer 8 -> 5, the way a neuron op does: the + # module object stays, its parameters are replaced. + model.feat = nn.Linear(16, 5) + model.classifier.weight = nn.Parameter(th.randn(3, 5)) + model.classifier.in_features = 5 + + model(th.randn(8, 16)) + tracker.observe_batch(list(range(8)), step=2) + self.assertEqual(tracker.steps_trained, 2, "projection died after the edit") + self.assertEqual(tracker.stats()["feature_dim"], 5) + self.assertEqual(tracker.rebuilds, 1) + finally: + tracker.detach() + + def test_hook_survives_a_parameter_swap(self): + """The hook is on the module, and WeightsLab's ops replace parameters + rather than modules -- so a prune must not need a re-attach.""" + model = TinyNet() + tracker = ProjectionTracker() + tracker.attach(model.classifier) + try: + model.classifier.weight = nn.Parameter(th.randn(3, 8)) + tracker._pending = None + model(th.randn(4, 16)) + self.assertIsNotNone(tracker._pending) + finally: + tracker.detach() + + def test_encoder_follows_the_model_to_a_new_device(self): + """.cuda() after wrapping must not strand the encoder on the CPU.""" + tracker = ProjectionTracker() + tracker._ensure_encoder(8, th.device("cpu")) + first = tracker._encoder + tracker._ensure_encoder(8, th.device("cpu")) + self.assertIs(tracker._encoder, first, "same width must not rebuild") + self.assertEqual(tracker.rebuilds, 0) + + def test_write_back_does_not_recurse(self): + """save_signals calls observe_batch, and the write-back calls + save_signals -- the guard is what stops that looping.""" + tracker = ProjectionTracker(every_n_steps=1) + tracker._writing_back = True + self.assertFalse(tracker.observe_batch([1, 2, 3, 4], step=1)) + + +class TestPerInstancePairing(unittest.TestCase): + """Detection / segmentation: the loss wrapper rebinds batch_ids to one entry + per ANNOTATION, but the forward hook saw one feature row per SAMPLE. The + projection must be handed the sample-level ids, or every detection run + silently projects nothing (the length check rejects the mismatch).""" + + def setUp(self): + from weightslab.backend import ledgers + self.ledgers = ledgers + self.observed = [] + + # Two samples, three annotations between them -- the layout the + # per-instance branch reads back out of the dataframe. + frame = pd.DataFrame( + {"x": [0, 1, 2, 3, 4]}, + index=pd.MultiIndex.from_tuples( + [(7, 0), (7, 1), (7, 2), (9, 0), (9, 1)], + names=["sample_id", "annotation_id"]), + ) + + class _DfStub: + _df = frame + + self._saved_get_df = ledgers.get_dataframe + self._saved_observe = wl_src._projection_observe_batch + self._saved_instance = wl_src.save_instance_signals + ledgers.get_dataframe = lambda *a, **k: _DfStub() + wl_src._projection_observe_batch = lambda ids, step: self.observed.append(ids) + wl_src.save_instance_signals = lambda **kwargs: None + + def tearDown(self): + self.ledgers.get_dataframe = self._saved_get_df + wl_src._projection_observe_batch = self._saved_observe + wl_src.save_instance_signals = self._saved_instance + + def test_projection_receives_sample_ids_not_annotation_ids(self): + sample_ids = [7, 9] + + def fake_forward(preds, targets): + return th.tensor([0.1, 0.2, 0.3, 0.4, 0.5]) # one value per instance + + wl_src.wrappered_fwd( + fake_forward, + {"per_sample": False, "per_instance": True, "log": False}, + "det_loss", + th.randn(2, 4), th.randn(2, 4), + batch_ids=sample_ids, + ) + + self.assertEqual(len(self.observed), 1) + # The sample ids, exactly as the hook saw them -- NOT the five + # annotation-level ids the per-instance branch rebinds batch_ids to. + self.assertEqual(list(self.observed[0]), sample_ids) + + +class TestOfflineReprojection(unittest.TestCase): + """``wl.project_dataset`` -- re-projecting a trained model from a different + layer, with no training step involved. The answer to "the live projection + picked the wrong layer".""" + + class _Stub: + def __init__(self): + self.calls = [] + + def enqueue_batch(self, **kwargs): + self.calls.append(kwargs) + + def setUp(self): + from torch.utils.data import DataLoader, TensorDataset + self._saved = wl_src.DATAFRAME_M + self.stub = self._Stub() + wl_src.DATAFRAME_M = self.stub + detach_projection() # no live tracker unless a test attaches one + + feats, labels = clustered(n_per=24, dim=16, clusters=4) + self.loader = DataLoader( + TensorDataset(feats, th.arange(len(feats)), th.tensor(labels)), + batch_size=16) + + class Net(nn.Module): + def __init__(self): + super().__init__() + self.early = nn.Sequential(nn.Linear(16, 12), nn.ReLU()) + self.late = nn.Sequential(nn.Linear(12, 9), nn.ReLU()) + self.classifier = nn.Linear(9, 4) + + def forward(self, x): + return self.classifier(self.late(self.early(x))) + + self.model = Net() + + def tearDown(self): + wl_src.DATAFRAME_M = self._saved + detach_projection() + + def _written(self, prefix): + return {k for call in self.stub.calls for k in call["losses"] + if k.startswith(f"signals//{prefix}_")} + + def test_projects_without_any_training_step(self): + stats = wl_src.project_dataset(self.model, self.loader, epochs=2, verbose=False) + self.assertEqual(stats["samples"], 96) + self.assertEqual(stats["feature_dim"], 9) # classifier's input + self.assertEqual(stats["prefix"], "umap") + self.assertEqual( + self._written("umap"), + {"signals//umap_x", "signals//umap_y", "signals//umap_z"}) + + def test_named_layer_overrides_the_auto_pick(self): + stats = wl_src.project_dataset( + self.model, self.loader, layer="late.0", epochs=2, verbose=False) + # late.0 is Linear(12, 9) -> its INPUT is 12-wide, not the auto-picked 9. + self.assertEqual(stats["feature_dim"], 12) + + def test_prefixes_coexist_so_layers_can_be_compared(self): + """The point of re-projecting: keep both and switch between them.""" + wl_src.project_dataset(self.model, self.loader, epochs=1, verbose=False) + wl_src.project_dataset(self.model, self.loader, layer="early.0", + prefix="umap_early", epochs=1, verbose=False) + self.assertTrue(self._written("umap")) + self.assertTrue(self._written("umap_early")) + self.assertIn("signals//umap_early_x", self._written("umap_early")) + + def test_model_weights_and_mode_are_untouched(self): + before = {k: v.clone() for k, v in self.model.state_dict().items()} + self.model.train() + wl_src.project_dataset(self.model, self.loader, epochs=2, verbose=False) + self.assertTrue(self.model.training, "training mode must be restored") + for key, value in self.model.state_dict().items(): + self.assertTrue(th.equal(before[key], value), f"{key} was modified") + + def test_no_hook_left_behind(self): + wl_src.project_dataset(self.model, self.loader, epochs=1, verbose=False) + self.assertEqual(len(self.model.classifier._forward_hooks), 0) + + def test_unknown_layer_name_lists_what_is_available(self): + with self.assertRaises(ValueError) as ctx: + wl_src.project_dataset(self.model, self.loader, layer="nope") + self.assertIn("nope", str(ctx.exception)) + self.assertIn("classifier", str(ctx.exception)) + + def test_layer_off_the_forward_path_says_so(self): + self.model.orphan = nn.Linear(4, 4) # never called in forward + with self.assertRaises(RuntimeError) as ctx: + wl_src.project_dataset(self.model, self.loader, layer="orphan", + epochs=1, verbose=False) + self.assertIn("forward path", str(ctx.exception)) + + def test_accepts_the_wrapped_handle_users_actually_hold(self): + """``wl.watch_or_edit(net, flag="model")`` returns a ledger Proxy, not an + nn.Module -- and that is what every training script keeps a reference to. + A layer lookup against it used to report that the model had no layers at + all, which is a baffling thing to be told about a model you can see.""" + class FakeProxy: + def __init__(self, inner): + self.model = inner + + def __call__(self, *args, **kwargs): + return self.model(*args, **kwargs) + + proxy = FakeProxy(self.model) + self.assertEqual(list(getattr(proxy, "named_modules", lambda: [])()), []) + + stats = wl_src.project_dataset(proxy, self.loader, layer="late.0", + epochs=1, verbose=False) + self.assertEqual(stats["feature_dim"], 12) + + def test_resolve_module_unwraps_and_passes_through(self): + from weightslab.projection import resolve_module + + class Wrapper: + def __init__(self, inner): + self.model = inner + + self.assertIs(resolve_module(self.model), self.model) + self.assertIs(resolve_module(Wrapper(self.model)), self.model) + self.assertIs(resolve_module(Wrapper(Wrapper(self.model))), self.model) + self.assertIsNone(resolve_module(object())) + + def test_offline_then_resume_adopts_by_default_on_a_prefix_collision(self): + """Writing to the prefix the LIVE projection owns and then resuming + training would overwrite these coordinates sample by sample with a + freshly-initialised encoder -- half this layout, half another, in one + cloud. The default is therefore to hand the fit over.""" + attach_projection(self.model, every_n_steps=1) + live = get_tracker() + try: + self.assertIsNone(live._encoder, "live encoder should not exist yet") + stats = wl_src.project_dataset( + self.model, self.loader, layer="late.0", epochs=1, verbose=False) + + self.assertTrue(stats["adopted_by_live"]) + # The live projection now owns the offline encoder AND is hooked to + # the layer it was fitted on -- adopting one without the other just + # triggers a rebuild and throws the fit away. + self.assertIsNotNone(live._encoder) + self.assertEqual(live._encoder_dim, 12) + self.assertEqual(live.signal_prefix, "umap") + + # Resuming training refines that same encoder instead of restarting. + fits_before = live.steps_trained + for step, (x, ids, _) in enumerate(self.loader, 1): + self.model(x) + live.observe_batch(ids, step=step) + self.assertGreater(live.steps_trained, fits_before) + self.assertEqual(live._encoder_dim, 12, "rebuilt instead of continuing") + finally: + detach_projection() + + def test_a_distinct_prefix_leaves_the_live_projection_alone(self): + """A snapshot under its own name is never written again, so there is + nothing to hand over.""" + attach_projection(self.model, every_n_steps=1) + live = get_tracker() + try: + stats = wl_src.project_dataset( + self.model, self.loader, layer="late.0", prefix="umap_late", + epochs=1, verbose=False) + self.assertFalse(stats["adopted_by_live"]) + self.assertIsNone(live._encoder, "live projection must be untouched") + self.assertEqual(live.signal_prefix, "umap") + finally: + detach_projection() + + def test_adopt_false_on_a_collision_warns_instead_of_corrupting_silently(self): + attach_projection(self.model, every_n_steps=1) + try: + with self.assertLogs("weightslab.projection.parametric_umap", + level="WARNING") as captured: + stats = wl_src.project_dataset( + self.model, self.loader, epochs=1, verbose=False, adopt=False) + self.assertFalse(stats["adopted_by_live"]) + message = "\n".join(captured.output) + self.assertIn("overwrite", message) + self.assertIn("adopt=True", message) + finally: + detach_projection() + + def test_adopt_true_switches_the_live_projection_to_a_new_prefix(self): + attach_projection(self.model, every_n_steps=1) + live = get_tracker() + try: + stats = wl_src.project_dataset( + self.model, self.loader, layer="early.0", prefix="umap_early", + epochs=1, verbose=False, adopt=True) + self.assertTrue(stats["adopted_by_live"]) + self.assertEqual(live.signal_prefix, "umap_early") + self.assertEqual(live._encoder_dim, 16) # early.0's input width + finally: + detach_projection() + + @unittest.skipIf(not th.cuda.is_available(), "needs a second device") + def test_adoption_across_devices_keeps_fitting(self): + """An adopted encoder can arrive on another device. Moving it with + .to() leaves Adam's state behind, so step() then raises on every batch + -- swallowed by observe_batch's broad except into a projection that + silently stops. The optimizer has to be rebuilt with it.""" + model = self.model.cuda() + attach_projection(model, every_n_steps=1) + live = get_tracker() + try: + wl_src.project_dataset(model, self.loader, layer="late.0", + epochs=1, verbose=False, device="cpu") + self.assertIsNotNone(live._encoder) + fits_before = live.steps_trained + for step, (x, ids, _) in enumerate(self.loader, 1): + model(x.cuda()) + live.observe_batch(ids, step=step) + break + self.assertGreater(live.steps_trained, fits_before, + "adopted encoder stopped fitting after the move") + self.assertEqual( + next(live._encoder.parameters()).device.type, + live._optimizer.param_groups[0]["params"][0].device.type) + finally: + detach_projection() + self.model = model.cpu() + + def test_no_live_projection_means_nothing_to_adopt(self): + detach_projection() + stats = wl_src.project_dataset(self.model, self.loader, epochs=1, + verbose=False) + self.assertFalse(stats["adopted_by_live"]) + + def test_max_samples_caps_collection(self): + stats = wl_src.project_dataset(self.model, self.loader, epochs=1, + max_samples=32, verbose=False) + self.assertEqual(stats["samples"], 32) + + def test_offline_fit_preserves_cluster_structure(self): + """Same bar the live path is held to, on the offline fitter.""" + feats, labels = clustered(n_per=40, dim=16, clusters=3) + from torch.utils.data import DataLoader, TensorDataset + loader = DataLoader( + TensorDataset(feats, th.arange(len(feats)), th.tensor(labels)), + batch_size=60) + model = nn.Sequential(nn.Identity(), nn.Linear(16, 4)) + wl_src.project_dataset(model, loader, epochs=40, fit_batch=120, + verbose=False) + + coords, ids = [], [] + for call in self.stub.calls: + if "signals//umap_x" not in call["losses"]: + continue + coords.append(np.stack( + [call["losses"][f"signals//umap_{a}"] for a in "xyz"], axis=1)) + ids.extend(int(s) for s in call["sample_ids"]) + coords = np.concatenate(coords) + got = labels[np.array(ids)] + + centroids = np.stack([coords[got == c].mean(axis=0) for c in range(3)]) + within = np.mean([np.linalg.norm(coords[got == c] - centroids[c], axis=1).mean() + for c in range(3)]) + between = np.mean([np.linalg.norm(centroids[i] - centroids[j]) + for i in range(3) for j in range(i + 1, 3)]) + self.assertGreater(between, 2.0 * within) + + +class TestFollowView(unittest.TestCase): + """Whose filter the projection obeys. + + A filter from the DATA side (typed, or asked of the agent) should narrow + the cloud with it. A filter the projection itself produced must not, or a + lasso collapses the view to the handful of points just picked and there is + nothing left to select from. + """ + + def _frame(self, n=200): + rng = np.random.default_rng(3) + index = pd.MultiIndex.from_arrays( + [np.where(np.arange(n) % 4 == 0, "test_loader", "train_loader"), + np.arange(n).astype(str)], + names=["origin", "sample_id"]) + return pd.DataFrame( + {f"signals//{a}": rng.random(n) for a in + ("umap_x", "umap_y", "umap_z")}, index=index) + + def test_request_carries_the_flag(self): + req = pb2.ProjectionRequest(follow_view=True) + self.assertTrue(req.follow_view) + self.assertFalse(pb2.ProjectionRequest().follow_view) + + def test_serving_is_indifferent_to_the_flag(self): + """The flag selects WHICH frame data_service hands over; the builder + itself just renders whatever it is given, either way.""" + frame = self._frame(120) + for follow in (True, False): + resp = ps.build_projection_response( + frame, pb2.ProjectionRequest(follow_view=follow)) + self.assertTrue(resp.success) + self.assertEqual(resp.total_available, 120) + + def test_a_filtered_frame_yields_a_filtered_cloud(self): + """What "the projection matches the current view" actually means: hand + it the narrowed frame and only those points come back.""" + full = self._frame(200) + narrowed = full.iloc[:40] + wide = ps.build_projection_response(full, pb2.ProjectionRequest(follow_view=True)) + thin = ps.build_projection_response(narrowed, pb2.ProjectionRequest(follow_view=True)) + self.assertEqual(wide.total_available, 200) + self.assertEqual(thin.total_available, 40) + self.assertTrue(set(thin.sample_ids).issubset(set(wide.sample_ids))) + + def test_discarded_flags_are_sent_for_greying(self): + frame = self._frame(50) + frame["discarded"] = [i % 5 == 0 for i in range(50)] + resp = ps.build_projection_response(frame, pb2.ProjectionRequest()) + self.assertEqual(len(resp.discarded), 50) + self.assertEqual(sum(resp.discarded), 10) + + def test_target_colours_points_without_being_asked(self): + """On a classification set the classes ARE the clusters; colouring by + split paints the whole train set one colour and says nothing.""" + frame = self._frame(60) + frame["target"] = np.arange(60) % 10 + resp = ps.build_projection_response(frame, pb2.ProjectionRequest()) + self.assertEqual(len(resp.color_values) + len(resp.color_labels), 60) + + +class TestStratifiedDecimation(unittest.TestCase): + """No cluster may vanish because the budget is small.""" + + def _blobs(self, sizes): + ids, groups = [], [] + for name, count in sizes.items(): + ids.extend(f"{name}-{i}" for i in range(count)) + groups.extend([name] * count) + return np.array(ids), np.array(groups) + + def test_small_clusters_keep_their_floor(self): + """A flat global sample gives a 200-point cluster ~2 points out of a + 100k cloud -- it reads as noise, or disappears, and the user concludes + the model has no such cluster.""" + sizes = {"big": 100_000, "mid": 5_000, "tiny": 200} + ids, groups = self._blobs(sizes) + + budget = 20_000 + keep = ps.stratified_keep(ids, groups, budget) + self.assertLessEqual(len(keep), budget) + + kept = groups[keep] + for name, count in sizes.items(): + got = int((kept == name).sum()) + floor = int(np.ceil(0.10 * count)) + self.assertGreaterEqual( + got, min(floor, count), + f"{name}: kept {got} of {count}, below 10% of itself") + + def test_tight_budget_still_protects_the_small_clusters(self): + """When 10% of everything exceeds the budget the big group absorbs the + shortfall and the small ones keep their share. Losing a little + resolution on a 100k blob costs nothing; losing a 200-point cluster + loses the finding.""" + sizes = {"big": 100_000, "mid": 5_000, "tiny": 200} + ids, groups = self._blobs(sizes) + budget = 5_000 + keep = ps.stratified_keep(ids, groups, budget) + self.assertLessEqual(len(keep), budget) + + kept = groups[keep] + for name in ("tiny", "mid"): + got = int((kept == name).sum()) + floor = int(np.ceil(0.10 * sizes[name])) + self.assertGreaterEqual(got, floor, f"{name} lost its share") + self.assertGreater(int((kept == "big").sum()), 0) + + def test_every_class_is_drawn_when_the_budget_is_far_below_the_data(self): + """THE scale regression. A 10% floor per group, paid smallest group + first, spent a 50k budget over 3M balanced samples on the first two + classes and drew the other eight not at all.""" + n, classes, budget = 300_000, 10, 5_000 + ids = np.arange(n) + groups = ids % classes + kept = groups[ps.stratified_keep(ids, groups, budget)] + counts = np.bincount(kept, minlength=classes) + self.assertEqual(int(counts.sum()), budget) + # Balanced classes, balanced draw. + self.assertGreater(counts.min(), 0.8 * budget / classes, counts) + + def test_a_small_cluster_is_drawn_denser_than_a_flat_sample(self): + """1/sqrt(size) stratification: 100x smaller is drawn ~10x denser.""" + ids = np.arange(1_010_000) + groups = np.where(ids < 10_000, 1, 0) # 10k cluster in a 1M cloud + kept = groups[ps.stratified_keep(ids, groups, 50_000)] + small_rate = (kept == 1).sum() / 10_000 + big_rate = (kept == 0).sum() / 1_000_000 + self.assertGreater(small_rate / big_rate, 5) + + def test_nothing_to_do_under_budget(self): + ids = np.array([f"s{i}" for i in range(50)]) + groups = np.array(["a"] * 25 + ["b"] * 25) + keep = ps.stratified_keep(ids, groups, 100) + self.assertEqual(len(keep), 50) + + def test_more_groups_than_budget_never_overshoots(self): + """500 groups into a 400-point budget: not every group can appear, and + quietly returning 500 points would blow the budget the client told us + it could draw. Spend it across as many groups as it holds.""" + ids = np.array([f"s{i}" for i in range(2000)]) + groups = np.array([f"g{i % 500}" for i in range(2000)]) + keep = ps.stratified_keep(ids, groups, 400) + self.assertLessEqual(len(keep), 400) + self.assertGreater(len(np.unique(groups[keep])), 300) + + def test_selection_is_deterministic(self): + ids = np.array([f"s{i}" for i in range(5000)]) + groups = np.array([f"g{i % 7}" for i in range(5000)]) + a = ps.stratified_keep(ids, groups, 500) + b = ps.stratified_keep(ids, groups, 500) + np.testing.assert_array_equal(a, b) + + def test_a_bigger_budget_only_adds_points(self): + """Zooming in raises the effective budget per region; a point on screen + must stay on screen.""" + ids = np.array([f"s{i}" for i in range(5000)]) + groups = np.array([f"g{i % 7}" for i in range(5000)]) + small = set(ps.stratified_keep(ids, groups, 300).tolist()) + large = set(ps.stratified_keep(ids, groups, 1200).tolist()) + self.assertTrue(small.issubset(large)) + + def test_grouping_prefers_the_colour_column(self): + n = 60 + frame = pd.DataFrame( + {"target": np.arange(n) % 3, "origin": ["train_loader"] * n}, + index=pd.Index(np.arange(n).astype(str), name="sample_id")) + coords = np.random.default_rng(0).random((n, 3)) + groups = ps.grouping_key(frame, coords, "target") + self.assertEqual(len(np.unique(groups)), 3) + + def test_grouping_falls_back_to_spatial_cells(self): + """No label at all -- the guarantee still has to mean something, so + group by where the points actually are.""" + n = 400 + frame = pd.DataFrame( + {"unrelated": np.arange(n)}, + index=pd.Index(np.arange(n).astype(str), name="sample_id")) + rng = np.random.default_rng(1) + coords = np.vstack([rng.normal(0, 0.2, (n // 2, 3)), + rng.normal(6, 0.2, (n // 2, 3))]) + groups = ps.grouping_key(frame, coords, "") + self.assertGreater(len(np.unique(groups)), 1) + # The two blobs must land in different cells. + self.assertNotEqual(groups[0], groups[-1]) + + +class TestPrefixDiscovery(unittest.TestCase): + """The board can only offer an offline re-projection if the server reports + it -- only the server knows which projections a dataset holds. + + These run with NO prefix registry, i.e. a run written before it existed, + which discovers projections from column pairs. The registered path is + covered in test_projection_lifecycle.py.""" + + def setUp(self): + from weightslab.projection.registry import clear_registry + self._saved_root = os.environ.pop("WEIGHTSLAB_ROOT_LOG_DIR", None) + clear_registry() + + def tearDown(self): + if self._saved_root is not None: + os.environ["WEIGHTSLAB_ROOT_LOG_DIR"] = self._saved_root + + def _frame(self, prefixes): + n = 40 + index = pd.MultiIndex.from_arrays( + [np.arange(n), np.zeros(n, dtype=int)], + names=["sample_id", "annotation_id"]) + data = {} + for prefix in prefixes: + for axis in "xyz": + data[f"signals//{prefix}_{axis}"] = np.random.rand(n) + return pd.DataFrame(data, index=index) + + def test_lists_every_stored_projection_with_the_live_one_first(self): + frame = self._frame(["umap_layer3", "umap", "umap_early"]) + self.assertEqual(ps.available_prefixes(frame), + ["umap", "umap_early", "umap_layer3"]) + + def test_response_carries_the_list(self): + frame = self._frame(["umap", "umap_layer3"]) + resp = ps.build_projection_response(frame, pb2.ProjectionRequest()) + self.assertEqual(list(resp.available_prefixes), ["umap", "umap_layer3"]) + + def test_a_named_prefix_selects_that_projection(self): + frame = self._frame(["umap", "umap_layer3"]) + resp = ps.build_projection_response( + frame, pb2.ProjectionRequest(prefix="umap_layer3")) + self.assertTrue(resp.success) + self.assertEqual(resp.returned, 40) + + def test_missing_prefix_reports_what_is_available(self): + frame = self._frame(["umap", "umap_layer3"]) + resp = ps.build_projection_response( + frame, pb2.ProjectionRequest(prefix="gone")) + self.assertFalse(resp.success) + self.assertIn("umap_layer3", resp.message) + self.assertEqual(list(resp.available_prefixes), ["umap", "umap_layer3"]) + + def test_half_written_projection_is_not_offered(self): + """A prefix with only _x would fail the moment it was selected.""" + frame = self._frame(["umap"]) + frame["signals//broken_x"] = np.random.rand(len(frame)) + self.assertEqual(ps.available_prefixes(frame), ["umap"]) + + +class TestEnvSwitch(unittest.TestCase): + def setUp(self): + self._saved = os.environ.get(ENV_ENABLED) + + def tearDown(self): + os.environ.pop(ENV_ENABLED, None) + if self._saved is not None: + os.environ[ENV_ENABLED] = self._saved + detach_projection() + + def test_default_is_on(self): + os.environ.pop(ENV_ENABLED, None) + self.assertTrue(projection_enabled()) + + def test_off_values(self): + for value in ("0", "false", "No", "OFF"): + os.environ[ENV_ENABLED] = value + self.assertFalse(projection_enabled(), value) + + def test_disabled_attaches_no_hook(self): + os.environ[ENV_ENABLED] = "0" + model = TinyNet() + self.assertIsNone(attach_projection(model)) + # Nothing captured, because nothing is hooked. + model(th.randn(4, 16)) + self.assertEqual(len(model.classifier._forward_hooks), 0) + + +class TestProjectionServing(unittest.TestCase): + """The view-box/level-of-detail contract.""" + + # The two index layouts that actually occur. The ledger frame is + # (sample_id, annotation_id); the data service's _all_datasets_df -- which + # is what GetProjection is handed -- is (origin, sample_id), with NO + # annotation level. Addressing levels positionally works on the first and + # silently returns zero rows on the second, so every serving test runs + # against both. + LAYOUTS = ("ledger", "view") + + def _index(self, n, layout): + if layout == "ledger": + return pd.MultiIndex.from_arrays( + [np.arange(n).astype(str), np.zeros(n, dtype=int)], + names=["sample_id", "annotation_id"]) + return pd.MultiIndex.from_arrays( + [np.where(np.arange(n) % 4 == 0, "test_loader", "train_loader"), + np.arange(n).astype(str)], + names=["origin", "sample_id"]) + + def _frame(self, n=500, prefix="umap", layout="view"): + rng = np.random.default_rng(7) + index = self._index(n, layout) + return pd.DataFrame({ + f"signals//{prefix}_x": rng.uniform(-10, 10, n), + f"signals//{prefix}_y": rng.uniform(-10, 10, n), + f"signals//{prefix}_z": rng.uniform(-10, 10, n), + "origin": rng.choice(["train_loader", "test_loader"], n), + "signals//loss": rng.uniform(0, 3, n), + }, index=index) + + def test_returns_packed_xyz_triples(self): + resp = ps.build_projection_response(self._frame(120), pb2.ProjectionRequest()) + self.assertTrue(resp.success) + self.assertEqual(resp.dims, 3) + self.assertEqual(len(resp.sample_ids), 120) + self.assertEqual(len(resp.coords), 120 * 3) + self.assertEqual(resp.total_available, 120) + + def test_missing_columns_report_cleanly(self): + frame = self._frame(10).drop(columns=["signals//umap_x", "signals//umap_y", + "signals//umap_z"]) + resp = ps.build_projection_response(frame, pb2.ProjectionRequest()) + self.assertFalse(resp.success) + self.assertIn("no projection for prefix", resp.message) + + def test_view_box_filters_to_the_box(self): + frame = self._frame(800) + req = pb2.ProjectionRequest(has_bounds=True, min_x=0, min_y=0, min_z=0, + max_x=10, max_y=10, max_z=10) + resp = ps.build_projection_response(frame, req) + self.assertTrue(resp.success) + xyz = np.array(resp.coords).reshape(-1, 3) + self.assertTrue((xyz >= 0).all() and (xyz <= 10).all()) + self.assertLess(resp.total_in_view, resp.total_available) + + def test_extent_describes_the_whole_cloud_not_the_box(self): + """The client frames the full cloud from this, so a zoomed-in request + must not shrink it.""" + frame = self._frame(400) + wide = ps.build_projection_response(frame, pb2.ProjectionRequest()) + req = pb2.ProjectionRequest(has_bounds=True, min_x=0, min_y=0, min_z=0, + max_x=1, max_y=1, max_z=1) + zoomed = ps.build_projection_response(frame, req) + self.assertAlmostEqual(wide.extent_min_x, zoomed.extent_min_x, places=5) + self.assertAlmostEqual(wide.extent_max_z, zoomed.extent_max_z, places=5) + + def test_budget_is_respected(self): + resp = ps.build_projection_response( + self._frame(5000), pb2.ProjectionRequest(max_points=250)) + self.assertEqual(len(resp.sample_ids), 250) + self.assertEqual(resp.total_in_view, 5000) + self.assertEqual(resp.total_available, 5000) + + def test_decimation_is_stable_across_requests(self): + """Two identical requests must pick the SAME points, or the cloud boils + under every camera move.""" + frame = self._frame(3000) + req = pb2.ProjectionRequest(max_points=300) + a = ps.build_projection_response(frame, req) + b = ps.build_projection_response(frame, req) + self.assertEqual(list(a.sample_ids), list(b.sample_ids)) + + def test_zooming_in_only_adds_points(self): + """A point drawn while zoomed out stays drawn when you zoom into it -- + the property that makes zoom feel like revealing detail, not resampling. + """ + frame = self._frame(4000) + wide = ps.build_projection_response(frame, pb2.ProjectionRequest(max_points=200)) + wide_xyz = np.array(wide.coords).reshape(-1, 3) + inside = ((wide_xyz >= 0).all(axis=1) & (wide_xyz <= 10).all(axis=1)) + kept_when_zoomed_out = {sid for sid, keep + in zip(wide.sample_ids, inside) if keep} + + zoomed = ps.build_projection_response(frame, pb2.ProjectionRequest( + has_bounds=True, min_x=0, min_y=0, min_z=0, + max_x=10, max_y=10, max_z=10, max_points=200)) + self.assertTrue(kept_when_zoomed_out.issubset(set(zoomed.sample_ids))) + + def test_origins_are_interned(self): + resp = ps.build_projection_response(self._frame(300), pb2.ProjectionRequest()) + self.assertEqual(sorted(resp.origins), ["test_loader", "train_loader"]) + self.assertEqual(len(resp.origin_ids), 300) + self.assertTrue(all(0 <= i < len(resp.origins) for i in resp.origin_ids)) + + def test_numeric_color_column(self): + resp = ps.build_projection_response( + self._frame(150), pb2.ProjectionRequest(color_column="signals//loss")) + self.assertEqual(len(resp.color_values), 150) + self.assertEqual(len(resp.color_labels), 0) + + def test_categorical_color_column(self): + frame = self._frame(150) + frame["cohort"] = np.where(np.arange(150) % 3 == 0, "alpha", "beta") + resp = ps.build_projection_response( + frame, pb2.ProjectionRequest(color_column="cohort")) + self.assertEqual(len(resp.color_labels), 150) + self.assertEqual(len(resp.color_values), 0) + + def test_colouring_by_origin_is_left_to_the_client(self): + """"Colour by split" is the one categorical column the server does NOT + answer. The Studio holds the split palette, so a sample is the same + colour in the cloud as in the data grid; sending labels here would give + the viewer a second, different palette for the same thing. + + See test_projection_colour_and_splits for the rest of this contract. + """ + resp = ps.build_projection_response( + self._frame(150), pb2.ProjectionRequest(color_column="origin")) + self.assertEqual(len(resp.color_labels), 0) + self.assertEqual(len(resp.color_values), 0) + # The split membership the client needs to do it is still sent. + self.assertEqual(sorted(resp.origins), ["test_loader", "train_loader"]) + + def test_instance_rows_are_excluded(self): + """Detection frames carry annotation_id >= 1 rows with no coordinates.""" + base = self._frame(100, layout="ledger") + extra = base.copy() + extra.index = pd.MultiIndex.from_arrays( + [np.arange(100).astype(str), np.ones(100, dtype=int)], + names=["sample_id", "annotation_id"]) + for axis in ("x", "y", "z"): + extra[f"signals//umap_{axis}"] = np.nan + resp = ps.build_projection_response(pd.concat([base, extra]), + pb2.ProjectionRequest()) + self.assertEqual(resp.total_available, 100) + + def test_every_index_layout_is_served(self): + """THE regression: the data service hands GetProjection a frame indexed + (origin, sample_id) -- no annotation level, and sample ids on level 1. + Addressing levels positionally read 'train_loader' as a sample id and + filtered every row away, so a working projection reported "no finite + coordinates yet" and the board never appeared.""" + for layout in self.LAYOUTS: + with self.subTest(layout=layout): + frame = self._frame(120, layout=layout) + resp = ps.build_projection_response(frame, pb2.ProjectionRequest()) + self.assertTrue(resp.success, f"{layout}: {resp.message}") + self.assertEqual(resp.total_available, 120) + self.assertEqual(len(resp.sample_ids), 120) + # Sample ids, not split names. + self.assertNotIn("train_loader", resp.sample_ids) + self.assertTrue(all(s.isdigit() for s in resp.sample_ids)) + + def test_single_level_index_is_served(self): + """A flat frame (no MultiIndex at all) must not be special-cased away.""" + n = 40 + rng = np.random.default_rng(3) + frame = pd.DataFrame( + {f"signals//umap_{a}": rng.random(n) for a in "xyz"}, + index=pd.Index(np.arange(n).astype(str), name="sample_id")) + resp = ps.build_projection_response(frame, pb2.ProjectionRequest()) + self.assertTrue(resp.success, resp.message) + self.assertEqual(resp.total_available, n) + + def test_unprojected_samples_are_skipped(self): + frame = self._frame(100) + frame.iloc[:40, frame.columns.get_loc("signals//umap_y")] = np.nan + resp = ps.build_projection_response(frame, pb2.ProjectionRequest()) + self.assertEqual(resp.total_available, 60) + + def test_the_frustum_keeps_only_what_the_camera_sees(self): + """A plane x >= 2 (n = (1,0,0), d = -2) plus far-away dummies.""" + frame = self._frame(2000) + far = [0, 0, 1, 1e6] * 5 + resp = ps.build_projection_response(frame, pb2.ProjectionRequest( + frustum_planes=[1, 0, 0, -2] + far)) + xyz = np.array(resp.coords).reshape(-1, 3) + self.assertTrue((xyz[:, 0] >= 2).all()) + self.assertGreater(len(xyz), 0) + self.assertLess(resp.total_in_view, resp.total_available) + + +class TestProjectionIndexAtScale(unittest.TestCase): + """The cell grid and the priority walk must choose exactly the points a + plain scan would -- they only read less to find them.""" + + def setUp(self): + self._saved = ps.CELL_INDEX_MIN_POINTS + ps.CELL_INDEX_MIN_POINTS = 1_000 # use the grid on a small cloud + + def tearDown(self): + ps.CELL_INDEX_MIN_POINTS = self._saved + + def _frame(self, n, dims=3, seed=5): + rng = np.random.default_rng(seed) + centres = rng.uniform(-20, 20, (6, dims)) + pick = rng.integers(0, 6, n) + xyz = centres[pick] + rng.normal(0, 2.0, (n, dims)) + index = pd.MultiIndex.from_arrays( + [np.where(np.arange(n) % 5 == 0, "test_loader", "train_loader"), np.arange(n)], + names=["origin", "sample_id"]) + data = {f"signals//umap_{a}": xyz[:, i] for i, a in enumerate("xyz"[:dims])} + data["target"] = pick + return pd.DataFrame(data, index=index) + + def _scan(self, index, region, budget): + hits = np.flatnonzero(region.contains(index.xyz)) + return hits[:budget], len(hits) + + def test_every_path_picks_the_same_points(self): + for dims in (2, 3): + index, failure = ps.ProjectionIndex.build(self._frame(30_000, dims), "umap", "") + self.assertIsNone(failure) + self.assertIsNotNone(index.cells) + boxes = [((-3,) * dims, (3,) * dims), # small: cell path + ((-50,) * dims, (50,) * dims), # everything: walk + ((-30,) * dims, (0,) * dims)] # in between + for lo, hi in boxes: + for budget in (100, 5_000, 50_000): + region = ps.ViewRegion(lo, hi) + got, in_view = index.select(region, budget) + want, total = self._scan(index, region, budget) + np.testing.assert_array_equal(got, want, f"{dims}d {lo} {budget}") + if in_view <= budget or total <= budget: + self.assertEqual(in_view, total) + + def test_culling_cells_against_the_frustum_loses_nothing(self): + """A wedge (x + y >= 0, z <= 5) inside a wide box: cells wholly + outside a plane are skipped, and the answer must not change.""" + index, _ = ps.ProjectionIndex.build(self._frame(30_000), "umap", "") + normals = np.array([[1, 1, 0], [0, 0, -1]], dtype=np.float32) + offsets = np.array([0, 5], dtype=np.float32) + region = ps.ViewRegion((-40,) * 3, (40,) * 3, normals, offsets) + for budget in (50, 2_000, 50_000): + got, _ = index.select(region, budget) + want, _ = self._scan(index, region, budget) + np.testing.assert_array_equal(got, want) + np.testing.assert_array_equal(index.select_all(region), + np.flatnonzero(region.contains(index.xyz))) + + def test_zooming_in_only_adds_points_through_the_grid(self): + index, _ = ps.ProjectionIndex.build(self._frame(40_000), "umap", "") + wide, _ = index.select(ps.ViewRegion((-40,) * 3, (40,) * 3), 2_000) + zoom = ps.ViewRegion((-5,) * 3, (5,) * 3) + narrow, _ = index.select(zoom, 2_000) + wide_inside = [p for p in wide if zoom.contains(index.xyz[p:p + 1])[0]] + self.assertTrue(set(wide_inside).issubset(set(narrow.tolist()))) + + def test_the_cached_answer_matches_the_uncached_one(self): + frame = self._frame(20_000) + request = pb2.ProjectionRequest(has_bounds=True, min_x=-10, min_y=-10, min_z=-10, + max_x=10, max_y=10, max_z=10, max_points=3_000) + cache = ps.ProjectionCache() + cached = cache.serve(request, view=lambda: frame, is_filtered=lambda: False) + plain = ps.build_projection_response(frame, request) + self.assertEqual(list(cached.sample_ids), list(plain.sample_ids)) + self.assertEqual(cached.total_in_view, plain.total_in_view) + + def test_moved_coordinates_reuse_the_order_and_draw_where_points_are(self): + frame = self._frame(5_000) + first, _ = ps.ProjectionIndex.build(frame, "umap", "") + same, _ = ps.ProjectionIndex.build(frame, "umap", "", previous=first) + self.assertIs(same, first) # nothing changed + frame["signals//umap_x"] += 100.0 # training moved the cloud + moved, _ = ps.ProjectionIndex.build(frame, "umap", "", previous=first) + self.assertIsNot(moved, first) + np.testing.assert_array_equal(moved.rows, first.rows) + self.assertGreater(moved.lo[0], first.lo[0] + 50) + # And an answer from the OLD index already draws the new positions. + resp = first.respond(pb2.ProjectionRequest(max_points=100)) + self.assertGreater(min(np.array(resp.coords).reshape(-1, 3)[:, 0]), 50) + + +class TestProjectionCache(unittest.TestCase): + def _frame(self, n=3000): + rng = np.random.default_rng(2) + index = pd.MultiIndex.from_arrays( + [np.full(n, "train_loader"), np.arange(n).astype(str)], + names=["origin", "sample_id"]) + return pd.DataFrame({f"signals//umap_{a}": rng.normal(size=n) for a in "xyz"}, + index=index) + + def test_a_grid_page_is_served_whole_ignoring_view_and_budget(self): + frame = self._frame() + cache = ps.ProjectionCache() + resp = cache.serve(pb2.ProjectionRequest( + follow_view=True, restrict_sample_ids=[str(i) for i in range(100, 160)], + max_points=10, has_bounds=True, min_x=50, max_x=60, min_y=50, max_y=60, + min_z=50, max_z=60), view=lambda: frame, is_filtered=lambda: False) + self.assertTrue(resp.success, resp.message) + self.assertEqual(sorted(map(int, resp.sample_ids)), list(range(100, 160))) + + def test_a_filtered_view_does_not_pull_the_dataset_on_every_request(self): + frame = self._frame() + pulls = [] + filtered = [False] + cache = ps.ProjectionCache() + kwargs = dict(view=lambda: frame, is_filtered=lambda: filtered[0], + pull_full=lambda: pulls.append(1) or frame) + cache.serve(pb2.ProjectionRequest(), **kwargs) # built from the full view + filtered[0] = True # a lasso filtered the grid + for _ in range(5): + resp = cache.serve(pb2.ProjectionRequest(), **kwargs) + self.assertEqual(resp.total_available, 3000) + self.assertEqual(pulls, []) + + def test_a_stale_index_is_refreshed_in_the_background(self): + frame = self._frame() + cache = ps.ProjectionCache() + cache.REFRESH_SECONDS = 0.0 + kwargs = dict(view=lambda: frame, is_filtered=lambda: False) + cache.serve(pb2.ProjectionRequest(), **kwargs) + before = cache._entries[("full", "")].index + frame["signals//umap_x"] += 1.0 + # Each request answers at once and, once the index is old enough + # (four times its build time), kicks a rebuild -- as the board's polls do. + deadline = time.time() + 10 + while cache._entries[("full", "")].index is before and time.time() < deadline: + cache.serve(pb2.ProjectionRequest(), **kwargs) + time.sleep(0.02) + self.assertIsNot(cache._entries[("full", "")].index, before) + + def test_an_edit_reaches_a_cloud_drawn_from_a_pulled_copy(self): + # A sorted or filtered grid: the whole-dataset cloud is drawn from a + # pulled copy, re-pulled only every FULL_PULL_SECONDS. A discard made + # in the grid used to stay un-greyed in the projection that long. + frame = self._frame(200) + frame["discarded"] = False + cache = ps.ProjectionCache() + kwargs = dict(view=lambda: frame, is_filtered=lambda: True, + pull_full=lambda: frame.copy()) + + def flags(): + resp = cache.serve(pb2.ProjectionRequest(max_points=1000), **kwargs) + return dict(zip(resp.sample_ids, resp.discarded)) + + self.assertFalse(flags()["7"]) + frame.loc[("train_loader", "7"), "discarded"] = True # the grid's edit + time.sleep(0.1) + self.assertFalse(flags()["7"]) # the copy is not re-pulled by itself + cache.note_edit() # what EditDataSample now does + deadline = time.time() + 10 + while not flags()["7"] and time.time() < deadline: + time.sleep(0.05) + self.assertTrue(flags()["7"]) + + def test_non_numeric_ids_rank_stably_whatever_else_is_in_the_array(self): + a = ps._stable_rank(np.array(["cat", "dog"])) + b = ps._stable_rank(np.array(["cat", "a-much-longer-sample-name", "dog"])) + self.assertEqual(a[0], b[0]) + self.assertEqual(a[1], b[2]) + + +class TestLassoSelection(unittest.TestCase): + """A lasso selects on the server, over every sample -- drawn or not.""" + + # Identity camera: clip = (x, y, z, 1), so NDC = (x, y), and a 100x100 + # canvas puts x = -1..1 at pixels 0..100 and y = 1..-1 at 0..100. + CAMERA = list(np.eye(4).T.ravel()) + VIEWPORT = [100, 100] + CENTRE = [25, 25, 75, 25, 75, 75, 25, 75] # NDC square [-0.5, 0.5]^2 + LEFT = [0, 0, 50, 0, 50, 100, 0, 100] # NDC x in [-1, 0] + + def _frame(self, n=20_000): + rng = np.random.default_rng(4) + index = pd.MultiIndex.from_arrays( + [np.full(n, "train_loader"), np.arange(n)], names=["origin", "sample_id"]) + return pd.DataFrame({f"signals//umap_{a}": rng.uniform(-1, 1, n) for a in "xyz"}, + index=index) + + def _lasso(self, cache, frame, polygon, mode, **extra): + return cache.serve(pb2.ProjectionRequest( + lasso_mode=mode, lasso_px=polygon, lasso_viewport=self.VIEWPORT, + view_projection=self.CAMERA, max_points=100, **extra), + view=lambda: frame, is_filtered=lambda: False) + + def _expected(self, frame, x0, x1, y0, y1): + x, y = frame["signals//umap_x"].to_numpy(), frame["signals//umap_y"].to_numpy() + return set(np.flatnonzero((x >= x0) & (x <= x1) & (y >= y0) & (y <= y1)).tolist()) + + def test_the_raster_test_matches_the_loop(self): + rng = np.random.default_rng(0) + xyz = rng.uniform(-1, 1, (5000, 3)).astype(np.float32) + inside = ps.points_in_lasso(xyz, self.CAMERA, np.array(self.CENTRE).reshape(-1, 2), + self.VIEWPORT) + want = (np.abs(xyz[:, 0]) < 0.48) & (np.abs(xyz[:, 1]) < 0.48) + clear = (np.abs(xyz[:, 0]) > 0.52) | (np.abs(xyz[:, 1]) > 0.52) + self.assertTrue(inside[want].all()) # well inside: selected + self.assertFalse(inside[clear].any()) # well outside: not + + def test_selects_every_sample_in_the_loop_not_just_the_drawn_ones(self): + frame = self._frame() + cache = ps.ProjectionCache() + resp = self._lasso(cache, frame, self.CENTRE, pb2.PROJECTION_LASSO_REPLACE) + self.assertTrue(resp.success, resp.message) + self.assertTrue(resp.selection_token) + self.assertGreater(resp.selected, 100 * 10) # far beyond the 100 drawn + expected = self._expected(frame, -0.5, 0.5, -0.5, 0.5) + got = set(cache.selection(resp.selection_token).tolist()) + # Pixel-quantised edges: everything well inside, nothing well outside. + self.assertTrue(self._expected(frame, -0.48, 0.48, -0.48, 0.48) <= got) + self.assertTrue(got <= self._expected(frame, -0.52, 0.52, -0.52, 0.52)) + self.assertAlmostEqual(len(got), len(expected), delta=0.05 * len(expected)) + + def test_refine_keeps_only_the_previous_selection_inside_the_new_loop(self): + frame = self._frame() + cache = ps.ProjectionCache() + first = self._lasso(cache, frame, self.CENTRE, pb2.PROJECTION_LASSO_REPLACE) + refined = self._lasso(cache, frame, self.LEFT, pb2.PROJECTION_LASSO_REFINE) + a = set(cache.selection(first.selection_token).tolist()) + b = set(cache.selection(refined.selection_token).tolist()) + self.assertTrue(b < a) + x = frame["signals//umap_x"].to_numpy() + self.assertTrue((x[sorted(b)] <= 0.025).all()) # one pixel = 0.02 + + def test_add_unions_with_the_previous_selection(self): + frame = self._frame() + cache = ps.ProjectionCache() + first = self._lasso(cache, frame, self.CENTRE, pb2.PROJECTION_LASSO_REPLACE) + both = self._lasso(cache, frame, self.LEFT, pb2.PROJECTION_LASSO_ADD) + a = set(cache.selection(first.selection_token).tolist()) + b = set(cache.selection(both.selection_token).tolist()) + self.assertTrue(a < b) + + def test_drawn_points_say_whether_they_are_selected(self): + """The client only ever holds the drawn subsample, so a point that + arrives with a later zoom must arrive marked.""" + frame = self._frame() + cache = ps.ProjectionCache() + kwargs = dict(view=lambda: frame, is_filtered=lambda: False) + resp = self._lasso(cache, frame, self.CENTRE, pb2.PROJECTION_LASSO_REPLACE) + selected = set(cache.selection(resp.selection_token).tolist()) + view = cache.serve(pb2.ProjectionRequest(max_points=5_000), **kwargs) + self.assertEqual(len(view.in_selection), len(view.sample_ids)) + self.assertEqual([int(s) in selected for s in view.sample_ids], list(view.in_selection)) + cleared = self._lasso(cache, frame, [], pb2.PROJECTION_LASSO_REPLACE) + self.assertEqual(cleared.selected, 0) + self.assertEqual(len(cache.serve(pb2.ProjectionRequest(), **kwargs).in_selection), 0) + + def test_the_data_board_applies_a_selection_by_token(self): + import contextlib + from types import SimpleNamespace + from weightslab.trainer.services.data_service import DataService + frame = self._frame() + cache = ps.ProjectionCache() + resp = self._lasso(cache, frame, self.CENTRE, pb2.PROJECTION_LASSO_REPLACE) + service = SimpleNamespace( + _projection_cache=cache, _all_datasets_df=frame, _is_filtered=False, + _watched_lock=lambda name: contextlib.nullcontext(), + _pull_into_all_data_view_df=lambda: frame, + _build_success_response=lambda df, message, intent_type: (df, message)) + df, _ = DataService._apply_projection_selection( + service, f"@projection_selection {resp.selection_token}") + self.assertEqual(len(df), resp.selected) + self.assertTrue(service._is_filtered) + expired = DataService._apply_projection_selection(service, "@projection_selection nope") + self.assertFalse(expired.success) + + def test_a_selection_finds_its_samples_in_a_rebuilt_view(self): + """The view the selection is applied to, or refined against, is often + not the frame the loop was drawn on (the data board rebuilt it in + between, reordered): its rows are found again by sample id.""" + frame = self._frame() + cache = ps.ProjectionCache() + resp = self._lasso(cache, frame, self.CENTRE, pb2.PROJECTION_LASSO_REPLACE) + rebuilt = frame.iloc[::-1].copy() # same samples, new rows + rows = cache.selection_rows(resp.selection_token, rebuilt) + self.assertEqual(len(rows), resp.selected) + picked = rebuilt.index.get_level_values("sample_id")[rows] + self.assertEqual(set(picked), set(cache.selection(resp.selection_token).tolist())) + # ...and a string-id view finds them too (the grid sends strings). + as_text = frame.copy() + as_text.index = pd.MultiIndex.from_arrays( + [frame.index.get_level_values(0), frame.index.get_level_values(1).astype(str)], + names=frame.index.names) + found = ps.SampleIdLookup().rows(as_text, [str(i) for i in picked[:50]]) + self.assertEqual(len(found), 50) + + +class TestDefaultViewBudget(unittest.TestCase): + """WEIGHTSLAB_PROJECTION_MAX_POINTS: how many points a view returns when the + client leaves the number to the server, which the board does by default.""" + + def setUp(self): + self._saved = os.environ.pop(ps.ENV_MAX_POINTS, None) + + def tearDown(self): + os.environ.pop(ps.ENV_MAX_POINTS, None) + if self._saved is not None: + os.environ[ps.ENV_MAX_POINTS] = self._saved + + def test_the_default_is_70k(self): + self.assertEqual(ps.default_max_points(), 70_000) + + def test_the_variable_sets_it_and_is_read_on_every_request(self): + os.environ[ps.ENV_MAX_POINTS] = "1500" + self.assertEqual(ps.default_max_points(), 1500) + os.environ[ps.ENV_MAX_POINTS] = "150_000" + self.assertEqual(ps.default_max_points(), 150_000) + + def test_it_is_capped_and_a_bad_value_falls_back(self): + os.environ[ps.ENV_MAX_POINTS] = "10000000" + self.assertEqual(ps.default_max_points(), ps.HARD_MAX_POINTS) + for bad in ("0", "-5", "lots", "70,000"): + os.environ[ps.ENV_MAX_POINTS] = bad + self.assertEqual(ps.default_max_points(), 70_000, bad) + + def test_a_view_returns_that_many_unless_the_client_asks_for_a_number(self): + rng = np.random.default_rng(3) + n = 5_000 + index = pd.MultiIndex.from_arrays([["train_loader"] * n, np.arange(n)], + names=["origin", "sample_id"]) + frame = pd.DataFrame({f"signals//umap_{a}": rng.normal(0, 5, n) for a in "xyz"}, + index=index) + frame["target"] = rng.integers(0, 4, n) + built, failure = ps.ProjectionIndex.build(frame, "umap", "") + self.assertIsNone(failure) + os.environ[ps.ENV_MAX_POINTS] = "1234" + self.assertEqual(built.respond(pb2.ProjectionRequest()).returned, 1234) + self.assertEqual(built.respond(pb2.ProjectionRequest(max_points=300)).returned, 300) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/general/test_projection_baseline_models.py b/tests/general/test_projection_baseline_models.py new file mode 100644 index 00000000..6b65bce0 --- /dev/null +++ b/tests/general/test_projection_baseline_models.py @@ -0,0 +1,222 @@ +"""The projection, swept across the whole baseline model zoo. + +``examples/utils/baseline_models`` is the closest thing the repo has to "every +shape of model a user might wrap": MLPs, CNNs, VGG/ResNet, U-Nets (2-D and +3-D), YOLO heads, a GAN and a VAE. The claim this feature makes is that it +attaches to any of them with no user input, so the honest way to hold that +claim is to actually run all of them. + +For each model the sweep asserts the auto-pick produced a ``(B, F)`` embedding +that is *not* simply the model's output width -- reading the head's output +instead of its input is the exact failure this is guarding against, and it +looks perfectly healthy unless you check the width. + +Generative models (GAN/VAE) are reported but not asserted on: their "head" is a +decoder producing an image, so there is no supervised representation to project +and the projection is not a meaningful view of them. +""" + +import inspect +import unittest + +import torch as th +import torch.nn as nn + +from weightslab.projection import ProjectionTracker, pick_embedding_layer + +try: + from weightslab.examples.utils.baseline_models.pytorch import models as zoo +except Exception: # pragma: no cover - the zoo is example data, not a hard dep + zoo = None + + +# Generative: no supervised representation, so the projection is not claimed +# to be meaningful for them (it still must not crash). +GENERATIVE = {"DCGAN", "SimpleVAE"} +# Not projection problems: +# - FlexibleCNNBlock needs constructor arguments the sweep cannot guess; +# - Yolov11 wraps an Ultralytics trainer whose forward loads a COCO dataset +# from disk, so running it here fails on missing data, not on the hook. +# TinyYOLO covers the same detection-head shape without that dependency. +SKIP = {"FlexibleCNNBlock", "Yolov11"} + + +def discover(): + """Every zero-arg-constructible model in the zoo, with its own input shape.""" + found = [] + for name, obj in vars(zoo).items(): + if not (inspect.isclass(obj) and issubclass(obj, nn.Module)): + continue + if obj.__module__ != zoo.__name__ or name in SKIP: + continue + try: + model = obj() + except Exception: + continue # needs constructor args + shape = getattr(model, "input_shape", None) + if not shape: + continue + found.append((name, model, tuple(shape))) + return sorted(found, key=lambda row: row[0]) + + +@unittest.skipIf(zoo is None, "baseline model zoo unavailable") +class TestBaselineModelZoo(unittest.TestCase): + """One subtest per model, so a failure names the model that broke.""" + + def test_every_baseline_model_yields_an_embedding(self): + models = discover() + self.assertGreater(len(models), 10, "zoo discovery found almost nothing") + + report = [] + for name, model, shape in models: + with self.subTest(model=name): + layer, use_input, feature_last = pick_embedding_layer(model) + self.assertIsNotNone(layer, f"{name}: nothing to hook") + + tracker = ProjectionTracker() + tracker.attach(layer, use_input=use_input, feature_last=feature_last) + model.eval() + # Batch of 2 at the model's own declared shape; a batch of 1 + # would hide any batch-dim mistake. + sample = th.randn((2,) + shape[1:]) + try: + with th.no_grad(): + output = model(sample) + finally: + handle_shape = tracker._pending + tracker.detach() + + captured = tuple(handle_shape.shape) if handle_shape is not None else None + out_width = None + if isinstance(output, th.Tensor): + out_width = output.shape[1] if output.ndim > 1 else None + report.append((name, type(layer).__name__, use_input, captured, out_width)) + + if name in GENERATIVE: + continue + + self.assertIsNotNone(captured, f"{name}: hook never fired") + self.assertEqual(len(captured), 2, f"{name}: not reduced to (B, F)") + self.assertEqual(captured[0], 2, f"{name}: lost the batch dim") + self.assertGreater(captured[1], 0, f"{name}: empty feature dim") + + print("\n--- baseline model zoo: projection auto-pick ---") + for name, layer, use_input, captured, out_width in report: + note = " (generative, not asserted)" if name in GENERATIVE else "" + print(f" {name:<28} {layer:<10} input={str(use_input):<5} " + f"-> {captured} model_out_width={out_width}{note}") + + def test_segmentation_models_project_features_not_class_logits(self): + """The regression that motivated hooking a conv's INPUT. + + A U-Net's final 1x1 conv outputs `n_classes` channels. Reading its + OUTPUT gives a (B, n_classes) tensor -- well-formed, and completely the + wrong thing to lay out. The embedding has to be wider than the class + count, which is what distinguishes the two. + """ + candidates = [(name, model) for name, model, _ in discover() + if "UNet" in name or "FCN" in name] + self.assertTrue(candidates, "no segmentation models discovered") + + for name, model in candidates: + with self.subTest(model=name): + layer, use_input, feature_last = pick_embedding_layer(model) + shape = tuple(model.input_shape) + tracker = ProjectionTracker() + tracker.attach(layer, use_input=use_input, feature_last=feature_last) + model.eval() + try: + with th.no_grad(): + output = model(th.randn((2,) + shape[1:])) + captured = tracker._pending + finally: + tracker.detach() + + self.assertIsNotNone(captured, f"{name}: hook never fired") + if isinstance(output, th.Tensor) and output.ndim > 1: + self.assertGreater( + captured.shape[1], output.shape[1], + f"{name}: captured {captured.shape[1]} features but the model " + f"outputs {output.shape[1]} channels -- that is the head's " + f"OUTPUT (class logits), not the representation feeding it", + ) + + +try: + import transformers # noqa: F401 + HAVE_TRANSFORMERS = True +except Exception: # pragma: no cover + HAVE_TRANSFORMERS = False + + +@unittest.skipIf(not HAVE_TRANSFORMERS, "transformers not installed") +class TestLanguageModels(unittest.TestCase): + """LLMs are the shape most likely to be projected wrongly. + + A language head is a ``Linear(hidden, vocab)`` with an enormous output, fed + ``(B, T, C)`` hidden states. Reading its OUTPUT would hand the projection a + vocab-sized logit vector; reading its INPUT and pooling the token axis gives + a sentence embedding, which is the thing worth laying out. These assert the + captured width is the model's hidden size and NOT its vocabulary/label count. + + Models are built from config (never downloaded) so the test is offline. + """ + + def _capture(self, model, **inputs): + layer, use_input, feature_last = pick_embedding_layer(model) + tracker = ProjectionTracker() + tracker.attach(layer, use_input=use_input, feature_last=feature_last) + model.eval() + try: + with th.no_grad(): + model(**inputs) + return tracker._pending, layer + finally: + tracker.detach() + + def test_causal_lm_projects_hidden_states_not_vocab_logits(self): + from transformers import GPT2Config, GPT2LMHeadModel + hidden, vocab = 128, 500 + model = GPT2LMHeadModel(GPT2Config( + n_embd=hidden, n_layer=2, n_head=4, vocab_size=vocab)) + captured, _ = self._capture(model, input_ids=th.randint(0, vocab, (2, 12))) + + self.assertIsNotNone(captured) + self.assertEqual(tuple(captured.shape), (2, hidden)) + self.assertNotEqual(captured.shape[1], vocab, + "projected the vocabulary logits, not the representation") + + def test_encoder_model(self): + from transformers import BertConfig, BertModel + hidden = 96 + model = BertModel(BertConfig( + hidden_size=hidden, num_hidden_layers=2, num_attention_heads=4, + intermediate_size=128, vocab_size=500)) + captured, _ = self._capture(model, input_ids=th.randint(0, 500, (2, 12))) + self.assertEqual(tuple(captured.shape), (2, hidden)) + + def test_sequence_classifier_projects_features_not_labels(self): + from transformers import BertConfig, BertForSequenceClassification + hidden, labels = 96, 3 + model = BertForSequenceClassification(BertConfig( + hidden_size=hidden, num_hidden_layers=2, num_attention_heads=4, + intermediate_size=128, vocab_size=500, num_labels=labels)) + captured, _ = self._capture(model, input_ids=th.randint(0, 500, (2, 12))) + self.assertEqual(tuple(captured.shape), (2, hidden)) + self.assertNotEqual(captured.shape[1], labels) + + def test_sequence_length_does_not_change_the_embedding_width(self): + """Token pooling means a batch of any length yields the same width -- + what makes sequences comparable as points at all.""" + from transformers import BertConfig, BertModel + model = BertModel(BertConfig( + hidden_size=96, num_hidden_layers=2, num_attention_heads=4, + intermediate_size=128, vocab_size=500)) + short, _ = self._capture(model, input_ids=th.randint(0, 500, (2, 4))) + long, _ = self._capture(model, input_ids=th.randint(0, 500, (2, 40))) + self.assertEqual(tuple(short.shape), tuple(long.shape)) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/general/test_projection_colour_and_splits.py b/tests/general/test_projection_colour_and_splits.py new file mode 100644 index 00000000..44d78865 --- /dev/null +++ b/tests/general/test_projection_colour_and_splits.py @@ -0,0 +1,319 @@ +"""Colour source, and the eval split reaching the cloud at all. + +Both of these were user-visible: the projection showed the train split only, +and colouring by class ran the class ids through a continuous ramp. +""" +import numpy as np +import pandas as pd + +from weightslab.proto import experiment_service_pb2 as pb2 +from weightslab.trainer.services import projection_service as ps + + +def _frame(n=64, targets=None, origins=None): + rng = np.random.default_rng(0) + sid = np.arange(n) + index = pd.MultiIndex.from_arrays( + [sid, np.zeros(n, dtype=int)], names=["sample_id", "annotation_id"]) + return pd.DataFrame({ + "signals//umap_x": rng.normal(size=n), + "signals//umap_y": rng.normal(size=n), + "signals//umap_z": rng.normal(size=n), + "target": np.arange(n) % 10 if targets is None else targets, + "origin": (["train_loader"] * n) if origins is None else origins, + "discarded": np.zeros(n, dtype=bool), + }, index=index) + + +def _request(**kwargs): + return pb2.ProjectionRequest(max_points=1000, **kwargs) + + +class TestColourSource: + def test_class_ids_are_categorical_not_a_ramp(self): + """A class id is a label, not a measurement. + + Through the continuous ramp, class 4 sits "between" 3 and 5 and + neighbouring classes come out near-identical -- the one thing a + projection coloured by class must not say. + """ + response = ps.build_projection_response(_frame(), _request(color_column="target")) + assert response.success + assert len(response.color_labels) == response.returned + assert not response.color_values + + def test_a_continuous_target_still_ramps(self): + # 100 distinct values: past the low-cardinality threshold, so this is a + # regression target and the ramp is the right reading of it. + frame = _frame(n=100, targets=np.linspace(0.0, 1.0, 100)) + response = ps.build_projection_response(frame, _request(color_column="target")) + assert response.success + assert len(response.color_values) == response.returned + assert not response.color_labels + + def test_colouring_by_split_is_left_to_the_client(self): + """The Studio owns the split palette, so a sample is the same colour + here as in the data grid. Sending labels would give the viewer a + second, different palette for the same thing.""" + frame = _frame(origins=["train_loader"] * 32 + ["test_loader"] * 32) + response = ps.build_projection_response(frame, _request(color_column="origin")) + assert response.success + assert not response.color_labels + assert not response.color_values + # ...and the split membership it needs to do that is still sent. + assert sorted(response.origins) == ["test_loader", "train_loader"] + + def test_a_metadata_column_still_works(self): + frame = _frame() + frame["signals//loss"] = np.linspace(0.0, 2.0, len(frame)) + response = ps.build_projection_response( + frame, _request(color_column="signals//loss")) + assert response.success + assert len(response.color_values) == response.returned + + +class TestBothSplitsAreServed: + def test_test_split_points_are_returned_alongside_train(self): + frame = _frame(origins=["train_loader"] * 40 + ["test_loader"] * 24) + response = ps.build_projection_response(frame, _request()) + assert response.success + names = {response.origins[i] for i in response.origin_ids} + assert names == {"train_loader", "test_loader"} + assert response.returned == 64 + + +class TestEvalBatchesArePlaced: + """The gate that kept the test split out of the cloud. + + ``every_n_steps`` rations FITTING, and only training batches fit. Applying + it to eval batches too meant a validation pass -- which does not advance the + step, so every batch of it arrives carrying the same number -- had its first + batch rejected unless that step happened to be a multiple of every_n_steps, + and all the rest rejected by ``step_i == self._last_step_run``. + """ + + def _tracker(self): + from weightslab.projection.parametric_umap import ProjectionTracker + tracker = ProjectionTracker(every_n_steps=50) + # Enough of a hook for observe_batch to believe one is attached. + tracker._handle = object() + return tracker + + def test_an_eval_batch_is_not_rationed_by_every_n_steps(self): + import torch as th + tracker = self._tracker() + placed = [] + tracker._write_coords = lambda ids, coords, step: placed.append(list(ids)) or True + + # A step that is NOT a multiple of every_n_steps, three times over, the + # way a validation pass arrives. + for _ in range(3): + tracker._pending = th.randn(8, 16) + tracker._pending_training = False + tracker.observe_batch(list(range(8)), step=137) + + assert len(placed) == 3, ( + "every eval batch should be placed; only the first got through") + + def test_a_training_batch_is_still_rationed(self): + import torch as th + tracker = self._tracker() + fits = [] + tracker._write_coords = lambda ids, coords, step: fits.append(list(ids)) or True + + tracker._pending = th.randn(8, 16) + tracker._pending_training = True + tracker.observe_batch(list(range(8)), step=137) # 137 % 50 != 0 + assert not fits + + tracker._pending = th.randn(8, 16) + tracker._pending_training = True + tracker.observe_batch(list(range(8)), step=150) # 150 % 50 == 0 + assert len(fits) == 1 + + def test_an_eval_batch_does_not_lock_out_the_training_batch_at_that_step(self): + """Only a fit claims the step. An eval batch that claimed it would shut + out the training batch arriving at the same one.""" + import torch as th + tracker = self._tracker() + seen = [] + tracker._write_coords = lambda ids, coords, step: seen.append(list(ids)) or True + + tracker._pending = th.randn(8, 16) + tracker._pending_training = False + tracker.observe_batch(list(range(8)), step=150) + + tracker._pending = th.randn(8, 16) + tracker._pending_training = True + tracker.observe_batch(list(range(100, 108)), step=150) + + assert len(seen) == 2 + + +class TestTheEncoderIsCheckpointedOncePerFit: + """A full evaluation pass must not rewrite the encoder once per batch. + + The checkpoint is keyed on ``steps_trained``, which only advances on a fit + -- but the write-back it lives in now also runs for eval batches, which are + placed rather than learned from. With the counter sitting on a multiple of + _SAVE_EVERY_FITS, every batch of an evaluation pass saved the same weights + again: a burst of saves a second apart for one fit's worth of encoder. + """ + + def test_repeated_eval_batches_save_at_most_once(self, monkeypatch): + import torch as th + from weightslab.projection import parametric_umap as pu + + saves = [] + monkeypatch.setattr(pu, "save_projection", lambda tracker: saves.append(1)) + + tracker = pu.ProjectionTracker(every_n_steps=50) + tracker._handle = object() + # Sitting exactly on a save boundary is the state that triggered it. + tracker.steps_trained = pu._SAVE_EVERY_FITS + # _write_coords imports it from weightslab.src at call time. + import weightslab.src as wl_src + monkeypatch.setattr(wl_src, "save_signals", lambda **kwargs: None) + + for _ in range(8): + tracker._pending = th.randn(8, 16) + tracker._pending_training = False + tracker.observe_batch(list(range(8)), step=137) + + assert len(saves) <= 1, f"encoder written {len(saves)} times for one fit" + + def test_a_later_fit_still_checkpoints(self, monkeypatch): + import torch as th + from weightslab.projection import parametric_umap as pu + + saves = [] + monkeypatch.setattr(pu, "save_projection", lambda tracker: saves.append(1)) + import weightslab.src as wl_src + monkeypatch.setattr(wl_src, "save_signals", lambda **kwargs: None) + + tracker = pu.ProjectionTracker(every_n_steps=50) + tracker._handle = object() + + tracker.steps_trained = pu._SAVE_EVERY_FITS + tracker._pending = th.randn(8, 16) + tracker._pending_training = False + tracker.observe_batch(list(range(8)), step=137) + + # The next save boundary is a different fit, and must still be written. + tracker.steps_trained = pu._SAVE_EVERY_FITS * 2 + tracker._pending = th.randn(8, 16) + tracker._pending_training = False + tracker.observe_batch(list(range(8)), step=139) + + assert len(saves) == 2 + + +class TestAuxiliaryHeadsAreNotHooked: + """Which layer the projection reads, per task family. + + torchvision's FCN/DeepLab register ``aux_classifier`` AFTER ``classifier``, + so "the last conv" chose the auxiliary head. That is wrong twice over: the + aux head branches off an earlier backbone stage, so its input is a + shallower representation than the one the model predicts from; and + torchvision skips it outside training, so on an evaluation pass the hook + produced nothing and the cloud silently stopped gaining test-split points. + """ + + def _picked(self, model): + from weightslab.projection import pick_embedding_layer + module, use_input, feature_last = pick_embedding_layer(model) + name = next((n for n, m in model.named_modules() if m is module), "?") + return name, module, use_input, feature_last + + def test_a_segmentation_model_with_an_aux_head_hooks_the_main_one(self): + import torch.nn as nn + + class Head(nn.Module): + def __init__(self): + super().__init__() + self.conv = nn.Conv2d(8, 4, 1) + + def forward(self, x): + return self.conv(x) + + class SegWithAux(nn.Module): + def __init__(self): + super().__init__() + self.backbone = nn.Conv2d(3, 8, 3, padding=1) + self.classifier = Head() + # Registered last, exactly as torchvision does it. + self.aux_classifier = Head() + + def forward(self, x): + return self.classifier(self.backbone(x)) + + name, _, use_input, feature_last = self._picked(SegWithAux()) + assert name == "classifier.conv", f"hooked {name}" + assert use_input is True and feature_last is False + + def test_the_same_holds_for_the_mmseg_spelling(self): + import torch.nn as nn + + class SegWithAux(nn.Module): + def __init__(self): + super().__init__() + self.backbone = nn.Conv2d(3, 8, 3, padding=1) + self.decode_head = nn.Conv2d(8, 4, 1) + self.auxiliary_head = nn.Conv2d(8, 4, 1) + + def forward(self, x): + return self.decode_head(self.backbone(x)) + + name, _, _, _ = self._picked(SegWithAux()) + assert name == "decode_head", f"hooked {name}" + + def test_a_model_that_is_only_an_aux_head_still_gets_hooked(self): + """The filter is a preference, not a veto: a model with nothing but an + auxiliary head must still project rather than fall through to a leaf.""" + import torch.nn as nn + + class OnlyAux(nn.Module): + def __init__(self): + super().__init__() + self.aux_classifier = nn.Conv2d(3, 4, 1) + + def forward(self, x): + return self.aux_classifier(x) + + name, module, _, _ = self._picked(OnlyAux()) + assert isinstance(module, nn.Conv2d) + assert name == "aux_classifier" + + def test_the_zoo_heads_per_task_family(self): + """The real models, by task: classification heads are Linear, the + segmentation and detection heads are the last conv of the MAIN head.""" + import inspect + import torch.nn as nn + try: + from weightslab.examples.utils.baseline_models.pytorch import models as zoo + except Exception: + import pytest + pytest.skip("baseline model zoo unavailable") + + names = {n: o for n, o in vars(zoo).items() + if inspect.isclass(o) and issubclass(o, nn.Module) + and o.__module__ == zoo.__name__} + expected = { + # classification -> the classifier Linear + "ResNet18": "model.fc", + "FashionCNN": "fc4", + # 2-D segmentation -> the output conv of the main head + "UNet": "outc", + "TinyUNet": "outc", + "UNet3p": "outc", + "FCNResNet50": "model.classifier.4", + # 3-D segmentation + "UNet3D": "out_conv", + # detection + "TinyYOLO": "head.1", + } + for model_name, want in expected.items(): + if model_name not in names: + continue + got, _, _, _ = self._picked(names[model_name]()) + assert got == want, f"{model_name}: hooked {got}, expected {want}" diff --git a/tests/general/test_projection_layer.py b/tests/general/test_projection_layer.py new file mode 100644 index 00000000..2e4999a3 --- /dev/null +++ b/tests/general/test_projection_layer.py @@ -0,0 +1,364 @@ +"""Which layer the live projection reads, and saying so on the board. + +The YOLO example never got a real projection: the auto-pick ("input of the +last conv") landed on the Detect head's frozen DFL box decoder, which training +never runs, and validation ran on an EMA copy built before the hook existed. +""" +import copy +import importlib.util +import json +import os +import shutil +import tempfile +import unittest + +import numpy as np +import pandas as pd +import torch as th +import torch.nn as nn + +from weightslab import src as wl_src +from weightslab.projection import ( + attach_projection, + clear_pending_restore, + detach_projection, + mirror_projection, + observe_batch, + pick_embedding_layer, + save_projection_coords, +) +from weightslab.projection import parametric_umap as pu +from weightslab.projection.registry import ( + REGISTRY_FILE, + clear_registry, + known_prefixes, + prefix_info, + register_prefix, +) +from weightslab.proto import experiment_service_pb2 as pb2 +from weightslab.trainer.services import projection_service as ps + + +class TinyNet(nn.Module): + def __init__(self): + super().__init__() + self.feat = nn.Linear(16, 8) + self.act = nn.ReLU() + self.classifier = nn.Linear(8, 3) + + def forward(self, x): + return self.classifier(self.act(self.feat(x))) + + +class DetectLike(nn.Module): + """A YOLO-style head: the class conv trains, and a frozen decoder conv is + registered after it but only runs at inference.""" + + def __init__(self): + super().__init__() + self.cls = nn.Conv2d(8, 2, 1) + self.dfl = nn.Conv2d(16, 1, 1, bias=False).requires_grad_(False) + + def forward(self, x): + return self.cls(x) + + +class YoloLike(nn.Module): + def __init__(self): + super().__init__() + self.body = nn.Conv2d(3, 8, 3, padding=1) + self.head = DetectLike() + + def forward(self, x): + return self.head(self.body(x)) + + +class _Stub: + """Stands in for the dataframe manager.""" + + def enqueue_batch(self, **kwargs): + pass + + def get_df_view(self, limit=-1, **_): + return pd.DataFrame() + + +def projection_hooks(module): + return [h for h in module._forward_hooks.values() if isinstance(h, pu._FeatureHook)] + + +class _Isolated(unittest.TestCase): + def setUp(self): + self._saved_frame = wl_src.DATAFRAME_M + wl_src.DATAFRAME_M = _Stub() + self.root = tempfile.mkdtemp(prefix="wl-projection-layer-") + self._saved_root = os.environ.get("WEIGHTSLAB_ROOT_LOG_DIR") + os.environ["WEIGHTSLAB_ROOT_LOG_DIR"] = self.root + detach_projection() + clear_pending_restore() + clear_registry() + + def tearDown(self): + wl_src.DATAFRAME_M = self._saved_frame + detach_projection() + clear_pending_restore() + clear_registry() + os.environ.pop("WEIGHTSLAB_ROOT_LOG_DIR", None) + if self._saved_root is not None: + os.environ["WEIGHTSLAB_ROOT_LOG_DIR"] = self._saved_root + shutil.rmtree(self.root, ignore_errors=True) + + +# --------------------------------------------------------------------------- +# the auto-pick +# --------------------------------------------------------------------------- +class TestFrozenLayersAreNotTheHead(unittest.TestCase): + + def test_a_frozen_decoder_after_the_head_is_skipped(self): + model = YoloLike() + module, use_input, _ = pick_embedding_layer(model) + self.assertIs(module, model.head.cls) + self.assertTrue(use_input) + + def test_the_picked_layer_runs_in_a_training_forward(self): + model = YoloLike().train() + tracker = attach_projection(model, every_n_steps=1) + try: + model(th.randn(4, 3, 8, 8)) + self.assertEqual(tracker._pending.shape, (4, 8)) + finally: + detach_projection() + + def test_a_model_frozen_whole_still_gets_a_pick(self): + model = YoloLike().requires_grad_(False) + module, _, _ = pick_embedding_layer(model) + self.assertIs(module, model.head.dfl) + + def test_trainable_models_pick_as_before(self): + model = TinyNet() + self.assertIs(pick_embedding_layer(model)[0], model.classifier) + + +# --------------------------------------------------------------------------- +# signals logged from inside the caller's no_grad +# --------------------------------------------------------------------------- +class TestFitsUnderTheCallersGradMode(_Isolated): + """Ultralytics ships every per-sample signal under ``no_grad``; the fit + used to inherit it, fail five times, and turn the projection off.""" + + def _run(self, mode): + model = TinyNet() + tracker = attach_projection(model, every_n_steps=1) + for step in range(1, 4): + model.train() + model(th.randn(16, 16)) + with mode(): + observe_batch(list(range(16)), step) + return tracker + + def test_fits_under_no_grad(self): + tracker = self._run(th.no_grad) + self.assertEqual(tracker.failures, 0) + self.assertEqual(tracker.steps_trained, 3) + self.assertGreater(tracker.samples_written, 0) + + def test_fits_under_inference_mode(self): + tracker = self._run(th.inference_mode) + self.assertEqual(tracker.failures, 0) + self.assertEqual(tracker.steps_trained, 3) + + def test_features_captured_under_inference_mode_still_fit(self): + model = TinyNet() + tracker = attach_projection(model, every_n_steps=1) + model.train() + with th.inference_mode(): + model(th.randn(16, 16)) + observe_batch(list(range(16)), 1) + self.assertEqual((tracker.failures, tracker.steps_trained), (0, 1)) + + def test_the_callers_grad_mode_is_left_as_it_was(self): + model = TinyNet() + attach_projection(model, every_n_steps=1) + model.train() + model(th.randn(16, 16)) + with th.no_grad(): + observe_batch(list(range(16)), 1) + self.assertFalse(th.is_grad_enabled()) + + +# --------------------------------------------------------------------------- +# a copy made before the hook (Ultralytics' EMA shadow) +# --------------------------------------------------------------------------- +class TestMirror(_Isolated): + + def setUp(self): + super().setUp() + self.model = TinyNet() + # Made BEFORE the attach, like UL's EMA: deepcopy cannot carry a hook + # that does not exist yet. + self.shadow = copy.deepcopy(self.model).eval() + self.live = attach_projection(self.model, every_n_steps=1) + + def test_a_copy_made_before_the_attach_feeds_nothing_on_its_own(self): + self.shadow(th.randn(5, 16)) + self.assertIsNone(self.live._pending) + + def test_a_mirrored_copy_feeds_the_live_tracker_as_eval(self): + self.assertTrue(mirror_projection(self.shadow)) + self.shadow(th.randn(5, 16)) + self.assertEqual(self.live._pending.shape, (5, 8)) + self.assertFalse(self.live._pending_training) + + def test_mirroring_twice_adds_one_hook(self): + mirror_projection(self.shadow) + mirror_projection(self.shadow) + self.assertEqual(len(projection_hooks(self.shadow.classifier)), 1) + + def test_detach_removes_the_mirror_too(self): + mirror_projection(self.shadow) + detach_projection() + self.assertEqual(projection_hooks(self.shadow.classifier), []) + + def test_reattaching_the_primary_keeps_the_mirror(self): + mirror_projection(self.shadow) + self.live.attach(self.model.classifier, use_input=True, feature_last=True) + self.assertEqual(len(projection_hooks(self.shadow.classifier)), 1) + + def test_a_copy_without_the_layer_is_refused(self): + self.assertFalse(mirror_projection(nn.Sequential(nn.Linear(16, 3)))) + + def test_nothing_to_mirror_without_a_live_projection(self): + detach_projection() + self.assertFalse(mirror_projection(self.shadow)) + + +# --------------------------------------------------------------------------- +# recording the layer, and serving it to the board +# --------------------------------------------------------------------------- +def _frame(prefix="umap", n=32): + rng = np.random.default_rng(0) + index = pd.MultiIndex.from_arrays( + [np.arange(n), np.zeros(n, dtype=int)], names=["sample_id", "annotation_id"]) + return pd.DataFrame({f"signals//{prefix}_{a}": rng.normal(size=n) for a in "xyz"}, + index=index) + + +class TestTheBoardIsToldTheLayer(_Isolated): + + def _live_fit(self): + model = TinyNet() + tracker = attach_projection(model, every_n_steps=1) + model.train() + model(th.randn(16, 16)) + observe_batch(list(range(16)), 1) + return tracker + + def test_the_live_projection_records_its_layer(self): + self._live_fit() + self.assertEqual(prefix_info("umap"), + {"layer": "classifier", "layer_detail": "Linear input"}) + + def test_the_layer_survives_a_restart(self): + self._live_fit() + detach_projection() + clear_registry() # a new process: memory is empty + self.assertEqual(prefix_info("umap")["layer"], "classifier") + + def test_the_response_names_the_live_layer(self): + self._live_fit() + response = ps.build_projection_response(_frame("umap"), pb2.ProjectionRequest(max_points=100)) + self.assertTrue(response.success) + self.assertEqual(response.layer, "classifier") + self.assertEqual(response.layer_detail, "Linear input") + + def test_the_live_tracker_wins_over_a_stale_registry_entry(self): + tracker = self._live_fit() + tracker.layer_name, tracker.layer_detail = "feat", "Linear input" + self.assertEqual(ps.projection_layer("umap"), ("feat", "Linear input")) + + def test_your_own_projection_has_no_layer(self): + save_projection_coords(np.ones((2, 3)), [1, 2], prefix="tsne") + response = ps.build_projection_response( + _frame("tsne"), pb2.ProjectionRequest(prefix="tsne", max_points=100)) + self.assertEqual((response.layer, response.layer_detail), ("", "")) + + def test_project_dataset_records_the_layer_it_read(self): + from torch.utils.data import DataLoader, TensorDataset + loader = DataLoader(TensorDataset(th.randn(16, 16), th.arange(16)), batch_size=8) + wl_src.project_dataset(TinyNet(), loader, layer="feat", method=lambda f: f[:, :2], + prefix="mine", verbose=False) + self.assertEqual(prefix_info("mine"), {"layer": "feat", "layer_detail": "Linear input"}) + + def test_registering_without_a_layer_keeps_the_recorded_one(self): + register_prefix("snap", layer="model.22", layer_detail="C3k2 output") + register_prefix("snap") + self.assertEqual(prefix_info("snap")["layer"], "model.22") + + def test_a_registry_file_from_before_layers_still_reads(self): + folder = os.path.join(self.root, "projection") + os.makedirs(folder) + with open(os.path.join(folder, REGISTRY_FILE), "w", encoding="utf-8") as handle: + json.dump(["umap", "tsne"], handle) + self.assertEqual(known_prefixes(), {"umap", "tsne"}) + self.assertEqual(prefix_info("tsne"), {}) + register_prefix("pca", layer="feat") + clear_registry() + self.assertEqual(known_prefixes(), {"umap", "tsne", "pca"}) + self.assertEqual(prefix_info("pca"), {"layer": "feat"}) + + def test_an_older_reader_still_finds_every_prefix(self): + """Older versions build the set with ``for n in json.load(...)``; on + the new object that iterates the keys.""" + register_prefix("umap", layer="classifier", layer_detail="Linear input") + register_prefix("tsne") + with open(os.path.join(self.root, "projection", REGISTRY_FILE), encoding="utf-8") as handle: + self.assertEqual({str(n) for n in json.load(handle)}, {"umap", "tsne"}) + + +# --------------------------------------------------------------------------- +# Ultralytics +# --------------------------------------------------------------------------- +# unittest's own skip, not pytest.importorskip: the release workflow runs this +# suite with `python -m unittest`, which reports pytest's Skipped as an error. +@unittest.skipUnless(importlib.util.find_spec("ultralytics") is not None, + "ultralytics not installed") +class TestUltralyticsLayer(_Isolated): + + def setUp(self): + super().setUp() + from ultralytics.nn.tasks import DetectionModel + self.model = DetectionModel("yolo11n.yaml", nc=2, verbose=False) + + def test_the_auto_pick_avoids_the_dfl(self): + module, _, _ = pick_embedding_layer(self.model) + name = next(n for n, m in self.model.named_modules() if m is module) + self.assertNotIn("dfl", name) + + def test_the_trainer_hooks_the_neck_output_and_it_trains(self): + from weightslab.integrations.ultralytics.trainer import _embedding_layer + layer = _embedding_layer(self.model) + tracker = attach_projection(self.model, layer=layer, every_n_steps=1) + self.assertEqual(tracker.layer_name, f"model.{len(self.model.model) - 2}") + self.assertTrue(tracker.layer_detail.endswith("output")) + + self.model.train() + self.model(th.randn(2, 3, 64, 64)) + self.assertIsNotNone(tracker._pending, "the hooked layer did not run in training") + self.assertEqual(tracker._pending.shape[0], 2) + self.assertTrue(tracker._pending_training) + + def test_the_ema_shadow_feeds_validation(self): + from ultralytics.utils.torch_utils import ModelEMA + from weightslab.integrations.ultralytics.trainer import _embedding_layer + ema = ModelEMA(self.model) # built before the hook, as UL does + tracker = attach_projection(self.model, layer=_embedding_layer(self.model), + every_n_steps=1) + self.assertTrue(mirror_projection(ema.ema)) + with th.no_grad(): + ema.ema(th.randn(2, 3, 64, 64)) + self.assertIsNotNone(tracker._pending) + self.assertFalse(tracker._pending_training) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/general/test_projection_lifecycle.py b/tests/general/test_projection_lifecycle.py new file mode 100644 index 00000000..7c09063e --- /dev/null +++ b/tests/general/test_projection_lifecycle.py @@ -0,0 +1,1337 @@ +"""The projection's life around the model: wrapping, opting out, failing, +checkpoints and restarts, and projections the user computes themselves. + +test_parametric_umap.py covers the algorithm and the serving contract; this +covers everything that happens to the projection BECAUSE of something done to +the model or the run. +""" + +import copy +import os +import shutil +import tempfile +import unittest +from unittest import mock + +import dill +import numpy as np +import pandas as pd +import torch as th +import torch.nn as nn + +import weightslab as wl +from weightslab import src as wl_src +from weightslab.backend import ledgers +from weightslab.components.checkpoint_manager import CheckpointManager +from weightslab.projection import ( + ENV_ENABLED, + ProjectionTracker, + attach_projection, + clear_pending_restore, + detach_projection, + get_tracker, + observe_batch, + reattach_projection, + save_projection, + save_projection_coords, +) +from weightslab.projection import parametric_umap as pu +from weightslab.projection.registry import ( + REGISTRY_FILE, + clear_registry, + known_prefixes, + register_prefix, +) +from weightslab.proto import experiment_service_pb2 as pb2 +from weightslab.trainer.services import projection_service as ps + +LOGGER = "weightslab.projection.parametric_umap" + + +class TinyNet(nn.Module): + def __init__(self): + super().__init__() + self.feat = nn.Linear(16, 8) + self.act = nn.ReLU() + self.classifier = nn.Linear(8, 3) + + def forward(self, x): + return self.classifier(self.act(self.feat(x))) + + +class _Stub: + """Stands in for the dataframe manager; records what would be written.""" + + def __init__(self, columns=()): + self.calls = [] + self._columns = list(columns) + + def enqueue_batch(self, **kwargs): + self.calls.append(kwargs) + + def get_df_view(self, limit=-1, **_): + return pd.DataFrame(columns=self._columns) + + def written(self, prefix): + return {k for c in self.calls for k in (c["losses"] or {}) + if k.startswith(f"signals//{prefix}_")} + + +def projection_hooks(module): + """The live projection hooks installed on *module*.""" + return [h for h in module._forward_hooks.values() if isinstance(h, pu._FeatureHook)] + + +def encoder_snapshot(): + return {k: v.detach().clone() for k, v in get_tracker()._encoder.state_dict().items()} + + +def same_weights(a, b): + return all(th.equal(a[k], b[k]) for k in a) + + +def fit(model, steps, start=1, batch=32): + """Run *steps* tracked training steps through the hook + observe path.""" + for step in range(start, start + steps): + model.train() + model(th.randn(batch, 16)) + observe_batch(list(range(batch)), step) + + +class _StubbedFrame(unittest.TestCase): + def setUp(self): + self._saved_frame = wl_src.DATAFRAME_M + self.stub = _Stub() + wl_src.DATAFRAME_M = self.stub + detach_projection() + clear_pending_restore() + + def tearDown(self): + wl_src.DATAFRAME_M = self._saved_frame + detach_projection() + clear_pending_restore() + + +class _LedgerRun(_StubbedFrame): + """A real ledger + checkpoint manager in a throwaway root_log_dir.""" + + def setUp(self): + super().setUp() + self.root = tempfile.mkdtemp(prefix="wl-projection-") + self._saved_env = os.environ.get(ENV_ENABLED) + os.environ.pop(ENV_ENABLED, None) + + def tearDown(self): + super().tearDown() + ledgers.clear_all() + # src caches the ledger's dataframe handle; clear_all orphans it, and a + # later test file writing through the stale handle fails with "Proxy + # target not set". Let it re-resolve, as a fresh process would. + wl_src.DATAFRAME_M = None + os.environ.pop(ENV_ENABLED, None) + if self._saved_env is not None: + os.environ[ENV_ENABLED] = self._saved_env + shutil.rmtree(self.root, ignore_errors=True) + + def boot(self, **model_kwargs): + """What a fresh training process does: ledger, HP, then the model.""" + ledgers.clear_all() + detach_projection() + manager = CheckpointManager(root_log_dir=self.root) + ledgers.register_checkpoint_manager(manager) + wl.watch_or_edit({"root_log_dir": self.root, + "experiment_dump_to_train_steps_ratio": 10 ** 9}, + flag="hyperparameters") + model_kwargs.setdefault("projection", {"every_n_steps": 1}) + model = wl.watch_or_edit(TinyNet(), flag="model", **model_kwargs) + return manager, model + + def checkpoint(self, manager): + manager.update_experiment_hash(first_time=True) + return manager.save_model_checkpoint(save_optimizer=False, update_manifest=True) + + +# --------------------------------------------------------------------------- +# (b) checkpoints and restarts +# --------------------------------------------------------------------------- +class TestCheckpointPairing(_LedgerRun): + + def test_restart_resumes_the_checkpoints_encoder_not_a_later_one(self): + """The start-up order: ModelInterface restores the checkpoint while it + is being built, the projection attaches only afterwards. The restore + used to find no tracker, do nothing, and leave the attach to pick up + the run-level encoder -- written AFTER the checkpoint, so it did not + match the weights just restored.""" + manager, model = self.boot() + fit(model, 3) + ckpt = self.checkpoint(manager) + at_checkpoint = encoder_snapshot() + fits_at_checkpoint = get_tracker().steps_trained + + fit(model, 5, start=10) + save_projection() # the run-level copy moves on + later = encoder_snapshot() + self.assertFalse(same_weights(at_checkpoint, later)) + + self.boot() # "restart the process" + tracker = get_tracker() + self.assertIsNotNone(tracker) + self.assertTrue(same_weights(encoder_snapshot(), at_checkpoint)) + self.assertEqual(tracker.steps_trained, fits_at_checkpoint) + self.assertTrue(CheckpointManager.projection_sidecar_path(ckpt).exists()) + + def test_restore_while_running_rolls_the_encoder_back(self): + """The Studio/agent order: the tracker is attached when the restore + lands, so it applies at once -- and drops the buffered features, which + came from the model that was just replaced.""" + manager, model = self.boot() + fit(model, 3) + ckpt = self.checkpoint(manager) + at_checkpoint = encoder_snapshot() + fit(model, 4, start=10) + self.assertTrue(get_tracker()._buffer_ids) + + self.assertTrue(manager._load_projection_sidecar(ckpt)) + self.assertTrue(same_weights(encoder_snapshot(), at_checkpoint)) + self.assertEqual(get_tracker()._buffer_ids, []) + + def test_a_checkpoint_from_before_the_first_fit_restores_an_empty_layout(self): + """No sidecar means no encoder existed at that step. Keeping the later + encoder would read the old model's features through a newer map.""" + manager, model = self.boot() + ckpt = self.checkpoint(manager) # nothing fitted yet + self.assertFalse(CheckpointManager.projection_sidecar_path(ckpt).exists()) + fit(model, 3) + self.assertIsNotNone(get_tracker()._encoder) + + self.assertFalse(manager._load_projection_sidecar(ckpt)) + tracker = get_tracker() + self.assertIsNone(tracker._encoder) + self.assertEqual(tracker.steps_trained, 0) + # ...and it simply starts learning again from there. + fit(model, 1, start=20) + self.assertIsNotNone(tracker._encoder) + + def test_the_sidecar_never_passes_for_a_weight_checkpoint(self): + manager, model = self.boot() + fit(model, 2) + ckpt = self.checkpoint(manager) + sidecar = CheckpointManager.projection_sidecar_path(ckpt) + self.assertTrue(sidecar.exists()) + self.assertEqual(sidecar.parent.name, "projection") + picked = manager._select_weight_checkpoint_file(manager.current_exp_hash) + self.assertEqual(picked.name, ckpt.name) + + def test_a_parked_restore_is_dropped_when_the_model_opts_out(self): + manager, model = self.boot() + fit(model, 2) + ckpt = self.checkpoint(manager) + detach_projection() + manager._load_projection_sidecar(ckpt) # no tracker: parked + self.assertIsNotNone(pu._PENDING_CHECKPOINT) + self.boot(projection=False) + self.assertIsNone(pu._PENDING_CHECKPOINT) + self.assertIsNone(get_tracker()) + + +class TestArchitecturePickling(_StubbedFrame): + """Model architectures are saved with dill, mid-run, hook attached.""" + + def setUp(self): + super().setUp() + self.model = TinyNet() + self.live = attach_projection(self.model, every_n_steps=1) + fit(self.model, 2) + + def test_the_architecture_file_does_not_carry_the_tracker(self): + bare = len(dill.dumps(TinyNet())) + hooked = len(dill.dumps(self.model)) + # Used to be ~160x: encoder, optimizer state and feature buffer rode + # along in every architecture file. + self.assertLess(hooked, bare * 1.5) + + def test_a_restored_architecture_feeds_the_live_tracker_once_reattached(self): + clone = dill.loads(dill.dumps(self.model)) + self.assertEqual(projection_hooks(clone.classifier), []) + self.live._pending = None + clone(th.randn(5, 16)) + self.assertIsNone(self.live._pending, "a dead hook fed something") + + self.assertTrue(reattach_projection(clone)) + clone(th.randn(5, 16)) + self.assertEqual(self.live._pending.shape, (5, 8)) + # The placeholder hook is gone and the old model is no longer hooked. + self.assertEqual(len(clone.classifier._forward_hooks), 1) + self.assertEqual(projection_hooks(self.model.classifier), []) + # The encoder survived the move: same layer, same width. + self.assertIsNotNone(self.live._encoder) + + def test_a_deep_copy_keeps_feeding_the_live_tracker(self): + """An EMA shadow model is a deepcopy; it should behave as before.""" + shadow = copy.deepcopy(self.model) + self.live._pending = None + shadow(th.randn(6, 16)) + self.assertEqual(self.live._pending.shape, (6, 8)) + + def test_reattach_finds_the_same_layer_by_name(self): + detach_projection() + attach_projection(self.model, layer="feat", every_n_steps=1) + clone = dill.loads(dill.dumps(self.model)) + reattach_projection(clone) + self.assertEqual(get_tracker().layer_name, "feat") + self.assertEqual(len(projection_hooks(clone.feat)), 1) + + +# --------------------------------------------------------------------------- +# (c) switching it off, and what happens when it cannot run +# --------------------------------------------------------------------------- +class TestModelWrapperSwitch(_LedgerRun): + + def test_projection_false_installs_nothing(self): + _, model = self.boot(projection=False) + self.assertIsNone(get_tracker()) + inner = pu.resolve_module(model) + self.assertTrue(all(not projection_hooks(m) for m in inner.modules())) + + def test_projection_false_removes_the_previous_models_hook(self): + _, first = self.boot() + first_inner = pu.resolve_module(first) + self.assertEqual(len(projection_hooks(first_inner.classifier)), 1) + self.boot(projection=False) + self.assertIsNone(get_tracker()) + self.assertEqual(projection_hooks(first_inner.classifier), []) + + def test_rewrapping_moves_the_hook_instead_of_adding_one(self): + _, first = self.boot() + _, second = self.boot() + self.assertEqual(projection_hooks(pu.resolve_module(first).classifier), []) + self.assertEqual(len(projection_hooks(pu.resolve_module(second).classifier)), 1) + + def test_a_string_is_shorthand_for_the_layer(self): + self.boot(projection="feat") + self.assertEqual(get_tracker().layer_name, "feat") + + def test_a_dict_steers_the_tracker(self): + self.boot(projection={"every_n_steps": 7, "graph_size": 64, + "n_neighbors": 5, "signal_prefix": "umap_head"}) + tracker = get_tracker() + self.assertEqual((tracker.every_n_steps, tracker.graph_size, + tracker.n_neighbors, tracker.signal_prefix), + (7, 64, 5, "umap_head")) + + def test_the_env_switch_wins_over_projection_true(self): + os.environ[ENV_ENABLED] = "0" + self.boot(projection=True) + self.assertIsNone(get_tracker()) + + def test_an_unknown_layer_falls_back_to_the_auto_pick(self): + with self.assertLogs(LOGGER, level="WARNING") as captured: + self.boot(projection={"layer": "no.such.layer"}) + self.assertIn("not found", "\n".join(captured.output)) + self.assertEqual(get_tracker().layer_name, "classifier") + + def test_an_unknown_option_is_a_warning_not_a_crash(self): + with self.assertLogs("weightslab.src", level="WARNING") as captured: + _, model = self.boot(projection={"every_n_step": 5}) # typo + self.assertIn("every_n_step", "\n".join(captured.output)) + self.assertIsNone(get_tracker()) + self.assertEqual(tuple(model(th.randn(2, 16)).shape), (2, 3)) + + def test_a_projection_that_cannot_install_never_stops_the_wrap(self): + with mock.patch.object(wl_src, "_attach_projection", + side_effect=RuntimeError("boom")): + with self.assertLogs("weightslab.src", level="WARNING") as captured: + _, model = self.boot() + self.assertIn("Could not install live projection", "\n".join(captured.output)) + self.assertIsNone(get_tracker()) + self.assertEqual(tuple(model(th.randn(2, 16)).shape), (2, 3)) + + +class TestFailureContract(_StubbedFrame): + """Instrumentation may stop; it may never take the run down.""" + + def setUp(self): + super().setUp() + self.model = TinyNet() + self.tracker = attach_projection(self.model, every_n_steps=1) + + def _fail_fits(self): + return mock.patch.object(self.tracker, "_umap_loss", + side_effect=RuntimeError("CUDA out of memory")) + + def test_a_failing_fit_warns_once_and_training_carries_on(self): + with self._fail_fits(), self.assertLogs(LOGGER, level="WARNING") as captured: + fit(self.model, 3) # raises nothing + warnings = [m for m in captured.output if "fit failed" in m] + self.assertEqual(len(warnings), 1) + self.assertIn("Training is unaffected", warnings[0]) + self.assertEqual(self.tracker.failures, 3) + self.assertIsNotNone(self.tracker._handle) # still trying + + def test_a_streak_of_failures_turns_the_projection_off(self): + with self._fail_fits(), self.assertLogs(LOGGER, level="WARNING") as captured: + fit(self.model, pu._MAX_CONSECUTIVE_FAILURES) + self.assertIsNone(self.tracker._handle) + self.assertIn("consecutive failed fits", self.tracker.disabled_reason) + self.assertIn("turned off", "\n".join(captured.output)) + stats = self.tracker.stats() + self.assertFalse(stats["enabled"]) + self.assertEqual(stats["failures"], pu._MAX_CONSECUTIVE_FAILURES) + # Nothing is hooked any more, and a restore does not revive it. + self.assertEqual(projection_hooks(self.model.classifier), []) + self.assertFalse(reattach_projection(self.model)) + self.assertEqual(tuple(self.model(th.randn(2, 16)).shape), (2, 3)) + + def test_one_good_fit_resets_the_streak(self): + n = pu._MAX_CONSECUTIVE_FAILURES - 1 + with self._fail_fits(): + fit(self.model, n, start=1) + fit(self.model, 1, start=50) + with self._fail_fits(): + fit(self.model, n, start=100) + self.assertIsNotNone(self.tracker._handle) + self.assertIsNone(self.tracker.disabled_reason) + + def test_nonfinite_features_never_reach_the_encoder(self): + """A diverged batch is dropped from the graph; the healthy samples + around it keep training the encoder, which stays finite.""" + fit(self.model, 1) + nan = lambda t, feature_last: th.full((t.shape[0], 8), float("nan")) + with mock.patch.object(pu, "flatten_features", side_effect=nan): + with self.assertLogs(LOGGER, level="WARNING") as captured: + for step in (10, 11, 12): + self.model(th.randn(32, 16)) + observe_batch(list(range(100, 132)), step) + self.assertEqual(sum("NaN/inf" in m for m in captured.output), 1) + self.assertTrue(all(th.isfinite(v).all() for v in encoder_snapshot().values())) + self.assertEqual(self.tracker.skipped_nonfinite, 3 * 32) # samples, not batches + self.assertEqual(self.tracker.failures, 0) + placed = {sid for c in self.stub.calls for sid in c["sample_ids"]} + self.assertFalse(placed & {str(i) for i in range(100, 132)}) + + def test_a_nonfinite_loss_is_never_stepped(self): + fit(self.model, 1) + before = encoder_snapshot() + nan = th.tensor(float("nan"), requires_grad=True) + with mock.patch.object(self.tracker, "_umap_loss", return_value=nan): + fit(self.model, 1, start=10) + self.assertTrue(same_weights(before, encoder_snapshot())) + self.assertEqual(self.tracker.failures, 1) + + def test_a_broken_hook_never_breaks_the_forward_pass(self): + with mock.patch.object(pu, "flatten_features", side_effect=ValueError("bad")): + out = self.model(th.randn(4, 16)) + self.assertEqual(tuple(out.shape), (4, 3)) + + +class TestLiveCoverage(_StubbedFrame): + """Every training batch feeds the graph and gets placed, not only the one + batch in every_n_steps that lands on a fit step.""" + + BATCH = 8 + + def setUp(self): + super().setUp() + self.model = TinyNet() + + def _step(self, step, n_batches=1): + """One training step; ids are unique per (step, micro-batch).""" + for k in range(n_batches): + self.model.train() + self.model(th.randn(self.BATCH, 16)) + base = (step * 10 + k) * self.BATCH + observe_batch(list(range(base, base + self.BATCH)), step) + + def _placed(self): + return {sid for c in self.stub.calls for sid in c["sample_ids"]} + + def _ids(self, step, k=0): + base = (step * 10 + k) * self.BATCH + return {str(i) for i in range(base, base + self.BATCH)} + + def test_non_fit_batches_reach_the_graph(self): + tracker = attach_projection(self.model, every_n_steps=5) + for step in range(1, 11): + self._step(step) + # The fit at step 10 saw the batches of steps 1-10, not just 5 and 10. + self.assertEqual(len(tracker._buffer_ids), 10 * self.BATCH) + self.assertTrue(self._ids(7) <= set(tracker._buffer_ids)) + + def test_every_sample_is_placed_once_the_encoder_has_fit(self): + attach_projection(self.model, every_n_steps=5) + for step in range(1, 21): + self._step(step) + placed = self._placed() + for step in range(1, 21): + self.assertTrue(self._ids(step) <= placed, f"step {step} not placed") + + def test_no_random_cloud_before_the_first_fit(self): + attach_projection(self.model, every_n_steps=5) + for step in range(1, 5): + self._step(step) + self.assertEqual(self.stub.calls, []) + + def test_buffer_memory_stays_bounded_between_fits(self): + tracker = attach_projection(self.model, every_n_steps=10 ** 6, graph_size=64) + for step in range(1, 51): + self._step(step) + self.assertLessEqual(tracker._buffer_count, 64 + self.BATCH) + self.assertEqual(tracker._buffer_count, len(tracker._buffer_ids)) + self.assertEqual(tracker._buffer_count, + sum(f.shape[0] for f in tracker._buffer_feats)) + + def test_gradient_accumulation_micro_batches_all_count(self): + """Several forwards share one optimizer step, hence one step number.""" + tracker = attach_projection(self.model, every_n_steps=10 ** 6) + self._step(1, n_batches=3) + self.assertEqual(len(tracker._buffer_ids), 3 * self.BATCH) + + def test_a_second_write_for_the_same_batch_is_not_a_miss(self): + """A flag="loss" criterion and a save_signals on the same batch both + call in; the second finds the features taken, which is not a miss.""" + attach_projection(self.model, every_n_steps=5) + with mock.patch.object(pu.logger, "warning") as warned: + for step in range(1, 11): + self._step(step) + observe_batch(list(range(self.BATCH)), step) # the second write + self.assertFalse(any("forward path" in str(c) for c in warned.call_args_list)) + + def test_a_layer_that_never_runs_is_still_reported(self): + tracker = attach_projection(self.model, layer="feat", every_n_steps=5) + tracker.detach() # never captures + tracker._handle = object() # but looks attached + with self.assertLogs(LOGGER, level="WARNING") as captured: + for step in range(1, 4): + observe_batch(list(range(self.BATCH)), step) + self.assertIn("forward path", "\n".join(captured.output)) + + +class TestSeenCounters(_StubbedFrame): + """Placing a sample is not the model seeing it.""" + + def test_the_live_write_back_carries_no_step(self): + model = TinyNet() + attach_projection(model, every_n_steps=1) + fit(model, 1, start=5) + writes = [c for c in self.stub.calls if "signals//umap_x" in (c["losses"] or {})] + self.assertTrue(writes) + # step=None is what keeps enqueue_batch off nb_seen / last_seen. + self.assertTrue(all(c["step"] is None for c in writes)) + + def test_an_ordinary_signal_still_counts_as_seen(self): + wl_src.save_signals(signals={"loss": [1.0, 2.0]}, batch_ids=[1, 2], step=9) + self.assertEqual(self.stub.calls[0]["step"], 9) + + def test_enqueue_without_a_step_leaves_the_counters_alone(self): + """The guarantee the two tests above rely on, on the real manager.""" + from weightslab.data.dataframe_manager import LedgeredDataFrameManager + from weightslab.data.sample_stats import SampleStats + manager = LedgeredDataFrameManager(enable_flushing_threads=False, enable_h5_persistence=False) + manager.enqueue_batch(sample_ids=["1"], preds_raw=None, preds=None, + losses={"signals//umap_x": np.array([0.5])}, step=None) + record = manager._buffer["1"] + self.assertNotIn(SampleStats.Ex.NB_SEEN.value, record) + self.assertNotIn(SampleStats.Ex.LAST_SEEN.value, record) + + +# --------------------------------------------------------------------------- +# (a) projections the user computes themselves +# --------------------------------------------------------------------------- +class TestBringYourOwnProjection(_StubbedFrame): + + def test_it_is_a_top_level_verb(self): + self.assertIs(wl.save_projection_coords, save_projection_coords) + + def test_coordinates_land_as_signals_under_the_prefix(self): + coords = np.arange(12, dtype=np.float32).reshape(4, 3) + columns = save_projection_coords(coords, batch_ids=[10, 11, 12, 13], prefix="tsne") + self.assertEqual(columns, ["signals//tsne_x", "signals//tsne_y", "signals//tsne_z"]) + call = self.stub.calls[0] + self.assertEqual(call["sample_ids"], ["10", "11", "12", "13"]) + np.testing.assert_allclose(call["losses"]["signals//tsne_y"], [1, 4, 7, 10]) + self.assertIsNone(call["step"]) # not counted as seen + + def test_tensors_and_lists_are_accepted(self): + save_projection_coords(th.zeros(3, 2), th.tensor([1, 2, 3]), prefix="pca") + save_projection_coords([[0, 1], [2, 3]], ["a", "b"], prefix="pca") + self.assertEqual(self.stub.written("pca"), {"signals//pca_x", "signals//pca_y"}) + + def test_a_2d_projection_over_a_stale_3d_one_flattens_z(self): + """The board reads x, y AND z when all three exist; new 2-D x/y would + otherwise be drawn against the old projection's depth.""" + wl_src.DATAFRAME_M = _Stub(columns=["signals//tsne_z"]) + columns = save_projection_coords(np.ones((2, 2)), [1, 2], prefix="tsne") + self.assertIn("signals//tsne_z", columns) + np.testing.assert_allclose(wl_src.DATAFRAME_M.calls[0]["losses"]["signals//tsne_z"], [0, 0]) + + def test_shapes_that_cannot_be_drawn_are_refused(self): + for bad in (np.ones(4), np.ones((4, 1)), np.ones((4, 4)), np.ones((4, 2, 1))): + with self.assertRaises(ValueError): + save_projection_coords(bad, [1, 2, 3, 4], prefix="tsne") + with self.assertRaises(ValueError) as ctx: + save_projection_coords(np.ones((3, 2)), [1, 2], prefix="tsne") + self.assertIn("aligned", str(ctx.exception)) + self.assertEqual(self.stub.calls, []) + + def test_prefixes_that_break_column_names_are_refused(self): + for bad in ("", " ", "a/b", "has space", "_lead", "signals//x"): + with self.assertRaises(ValueError, msg=bad): + save_projection_coords(np.ones((2, 2)), [1, 2], prefix=bad) + save_projection_coords(np.ones((2, 2)), [1, 2], prefix="pca.layer3-v2") + + def test_the_live_umap_prefix_is_refused_only_while_it_runs(self): + attach_projection(TinyNet()) + with self.assertRaises(ValueError) as ctx: + save_projection_coords(np.ones((2, 3)), [1, 2], prefix="umap") + self.assertIn("projection=False", str(ctx.exception)) + save_projection_coords(np.ones((2, 3)), [1, 2], prefix="tsne") + detach_projection() + save_projection_coords(np.ones((2, 3)), [1, 2], prefix="umap") # now free + + def test_it_is_flushed_at_once_so_the_board_sees_it(self): + """The board reads flushed rows; a buffered write stayed invisible for + a whole flush interval right after the user plugged it in.""" + self.stub.flush = mock.MagicMock() + save_projection_coords(np.ones((2, 2)), [1, 2], prefix="tsne") + self.stub.flush.assert_called_once() + + def test_nan_rows_are_allowed_and_mean_not_placed(self): + coords = np.array([[0, 0], [np.nan, np.nan]], dtype=np.float32) + save_projection_coords(coords, [1, 2], prefix="tsne") + self.assertTrue(np.isnan(self.stub.calls[0]["losses"]["signals//tsne_x"][1])) + + +class TestProjectDatasetWithYourMethod(_StubbedFrame): + """``project_dataset(method=...)``: WeightsLab collects the features, your + algorithm lays them out.""" + + def setUp(self): + super().setUp() + from torch.utils.data import DataLoader, TensorDataset + x = th.randn(40, 16) + self.loader = DataLoader(TensorDataset(x, th.arange(40)), batch_size=8) + self.model = TinyNet() + + def test_a_callable(self): + def first_two(features): + return features[:, :2] + stats = wl_src.project_dataset(self.model, self.loader, method=first_two, + verbose=False) + self.assertEqual(stats["prefix"], "first_two") + self.assertEqual(stats["method"], "first_two") + self.assertEqual(stats["feature_dim"], 8) # the head's input + self.assertEqual(self.stub.written("first_two"), + {"signals//first_two_x", "signals//first_two_y"}) + + def test_an_sklearn_style_estimator(self): + class FakeTSNE: + def __init__(self, n_components=3): + self.n_components = n_components + self.seen = None + + def fit_transform(self, features): + self.seen = features + return features[:, :self.n_components] + + estimator = FakeTSNE() + stats = wl_src.project_dataset(self.model, self.loader, method=estimator, + layer="feat", verbose=False) + self.assertEqual(stats["prefix"], "faketsne") + self.assertEqual(estimator.seen.shape, (40, 16)) # input of `feat` + self.assertIsInstance(estimator.seen, np.ndarray) + + def test_an_explicit_prefix_and_a_lambda(self): + stats = wl_src.project_dataset(self.model, self.loader, prefix="mine", + method=lambda f: f[:, :3], verbose=False) + self.assertEqual(stats["prefix"], "mine") + stats = wl_src.project_dataset(self.model, self.loader, + method=lambda f: f[:, :3], verbose=False) + self.assertEqual(stats["prefix"], "custom") + + def test_no_encoder_is_saved_or_adopted(self): + live = attach_projection(TinyNet()) + with mock.patch.object(pu, "save_projection") as saved: + stats = wl_src.project_dataset(self.model, self.loader, prefix="pca", + method=lambda f: f[:, :2], verbose=False) + saved.assert_not_called() + self.assertFalse(stats["adopted_by_live"]) + self.assertIs(get_tracker(), live) + self.assertEqual(live.signal_prefix, "umap") + + def test_mistakes_are_caught_before_the_feature_sweep(self): + attach_projection(TinyNet()) + swept = mock.MagicMock(side_effect=AssertionError("swept")) + with mock.patch.object(pu, "tqdm", swept): + with self.assertRaises(ValueError): + wl_src.project_dataset(self.model, self.loader, method="tsne") + with self.assertRaises(ValueError): # the live prefix + wl_src.project_dataset(self.model, self.loader, prefix="umap", + method=lambda f: f[:, :2]) + + def test_a_method_returning_the_wrong_shape_says_so(self): + with self.assertRaises(ValueError) as ctx: + wl_src.project_dataset(self.model, self.loader, + method=lambda f: f[:, :5], verbose=False) + self.assertIn("(N, 2) or (N, 3)", str(ctx.exception)) + + def test_the_built_in_method_is_unchanged(self): + stats = wl_src.project_dataset(self.model, self.loader, epochs=1, verbose=False) + self.assertEqual((stats["prefix"], stats["method"]), ("umap", "umap")) + + +class _StatefulLoader: + """A tracked WeightsLab loader's contract: ONE iterator, which a ``for`` + continues rather than restarts, with ``reset_iterator`` to start over.""" + + def __init__(self, x, batch=8): + self.batches = [(x[i:i + batch], th.arange(i, min(i + batch, len(x)))) + for i in range(0, len(x), batch)] + self.position = 0 + + def __len__(self): + return len(self.batches) + + def __iter__(self): + return self + + def __next__(self): + if self.position >= len(self.batches): + self.position = 0 # the next pass starts fresh + raise StopIteration + batch = self.batches[self.position] + self.position += 1 + return batch + + def reset_iterator(self): + self.position = 0 + + +class _AgedNet(TinyNet): + """A tracked model's age rule: a forward is a training step while the + tracking mode is TRAIN -- which a training guard leaves behind on exit.""" + + def __init__(self): + super().__init__() + from weightslab.components.tracking import TrackingMode + self.tracking_mode = TrackingMode.TRAIN + self.age = 0 + + def set_tracking_mode(self, mode): + self.tracking_mode = mode + + def forward(self, x): + from weightslab.components.tracking import TrackingMode + if self.tracking_mode == TrackingMode.TRAIN: + self.age += 1 + return super().forward(x) + + +class TestProjectDatasetDoesNotAgeTheModel(_StubbedFrame): + """The feature sweep is not training: it must not move the model's age.""" + + def setUp(self): + super().setUp() + from torch.utils.data import DataLoader, TensorDataset + self.loader = DataLoader(TensorDataset(th.randn(40, 16), th.arange(40)), batch_size=8) + self.model = _AgedNet() + + def test_a_sweep_leaves_the_age_alone(self): + wl_src.project_dataset(self.model, self.loader, method=lambda f: f[:, :2], + verbose=False) + self.assertEqual(self.model.age, 0) # five batches, no steps + + def test_the_built_in_method_leaves_it_alone_too(self): + wl_src.project_dataset(self.model, self.loader, epochs=1, verbose=False) + self.assertEqual(self.model.age, 0) + + def test_the_previous_tracking_mode_comes_back(self): + from weightslab.components.tracking import TrackingMode + wl_src.project_dataset(self.model, self.loader, method=lambda f: f[:, :2], + verbose=False) + self.assertEqual(self.model.tracking_mode, TrackingMode.TRAIN) + + def test_it_comes_back_when_the_sweep_itself_fails(self): + from weightslab.components.tracking import TrackingMode + + def breaks_after_one_batch(): + yield th.randn(8, 16), th.arange(8) + raise RuntimeError("the loader died") + + with self.assertRaises(RuntimeError): + wl_src.project_dataset(self.model, breaks_after_one_batch(), + method=lambda f: f[:, :2], verbose=False) + self.assertEqual(self.model.tracking_mode, TrackingMode.TRAIN) + + def test_a_model_that_does_not_track_is_untouched(self): + plain = TinyNet() # no set_tracking_mode at all + stats = wl_src.project_dataset(plain, self.loader, method=lambda f: f[:, :2], + verbose=False) + self.assertEqual(stats["samples"], 40) + + +class TestProjectDatasetOnALoaderMidEpoch(_StubbedFrame): + """A training loop is nearly always mid-epoch on its loader. The sweep must + still see the WHOLE split, not the rest of that epoch.""" + + def setUp(self): + super().setUp() + self.loader = _StatefulLoader(th.randn(40, 16)) + self.model = TinyNet() + for _ in range(3): # the training loop's three steps + next(self.loader) + + def test_the_sweep_covers_the_whole_split(self): + stats = wl_src.project_dataset(self.model, self.loader, + method=lambda f: f[:, :2], verbose=False) + self.assertEqual(stats["samples"], 40) # not the 16 left of the epoch + + def test_every_sample_id_is_placed_exactly_once(self): + wl_src.project_dataset(self.model, self.loader, prefix="pca", + method=lambda f: f[:, :2], verbose=False) + placed = [i for call in self.stub.calls for i in call["sample_ids"]] + self.assertEqual(sorted(int(i) for i in placed), list(range(40))) + + def test_a_finished_sweep_leaves_training_a_fresh_epoch(self): + wl_src.project_dataset(self.model, self.loader, + method=lambda f: f[:, :2], verbose=False) + first, _ = next(self.loader) + self.assertEqual(self.loader.position, 1) # the first batch of a new pass + self.assertTrue(th.equal(first, self.loader.batches[0][0])) + + def test_stopping_early_does_not_strand_training_mid_sweep(self): + wl_src.project_dataset(self.model, self.loader, max_samples=16, + method=lambda f: f[:, :2], verbose=False) + next(self.loader) + self.assertEqual(self.loader.position, 1) # not 3 + the batches swept + + def test_the_built_in_method_sweeps_the_whole_split_too(self): + stats = wl_src.project_dataset(self.model, self.loader, epochs=1, + verbose=False) + self.assertEqual(stats["samples"], 40) + + def test_the_sample_ids_follow_the_feature_rows(self): + """``sample_ids[n]`` is the sample whose features were row ``n`` -- what + lets a caller pair the rows its method received with labels.""" + from torch.utils.data import DataLoader, TensorDataset + x = th.randn(40, 16) + shuffled = DataLoader(TensorDataset(x, th.arange(40)), batch_size=8, shuffle=True) + seen = {} + + def method(features): + seen["features"] = features + return features[:, :2] + + stats = wl_src.project_dataset(self.model, shuffled, layer="feat", method=method, + verbose=False) + order = [int(i) for i in stats["sample_ids"]] + self.assertEqual(sorted(order), list(range(40))) + self.assertNotEqual(order, list(range(40))) # really shuffled + np.testing.assert_allclose(seen["features"], x[order].numpy(), rtol=1e-6) + + def test_a_plain_iterable_is_left_alone(self): + from torch.utils.data import DataLoader, TensorDataset + plain = DataLoader(TensorDataset(th.randn(40, 16), th.arange(40)), batch_size=8) + stats = wl_src.project_dataset(self.model, plain, + method=lambda f: f[:, :2], verbose=False) + self.assertEqual(stats["samples"], 40) + + +class TestServingAProjectionYouBrought(unittest.TestCase): + + def setUp(self): + detach_projection() + self._saved_root = os.environ.pop("WEIGHTSLAB_ROOT_LOG_DIR", None) + clear_registry() + for prefix in ("tsne", "umap"): + register_prefix(prefix) + + def tearDown(self): + detach_projection() + clear_registry() + if self._saved_root is not None: + os.environ["WEIGHTSLAB_ROOT_LOG_DIR"] = self._saved_root + + def _frame(self, prefixes): + n = 50 + rng = np.random.default_rng(0) + index = pd.MultiIndex.from_arrays( + [np.full(n, "train_loader"), np.arange(n).astype(str)], + names=["origin", "sample_id"]) + columns = {} + for prefix in prefixes: + for axis in "xy": + columns[f"signals//{prefix}_{axis}"] = rng.normal(size=n) + return pd.DataFrame(columns, index=index) + + def test_no_preference_picks_what_exists(self): + """A run with the built-in off and its own t-SNE used to fail the first + request (asking for 'umap') and only draw a poll later.""" + response = ps.build_projection_response(self._frame(["tsne"]), + pb2.ProjectionRequest()) + self.assertTrue(response.success, response.message) + self.assertEqual(list(response.available_prefixes), ["tsne"]) + self.assertEqual(response.dims, 2) + + def test_the_live_projection_is_still_the_default_when_present(self): + frame = self._frame(["tsne", "umap"]) + response = ps.build_projection_response(frame, pb2.ProjectionRequest()) + self.assertEqual(list(response.available_prefixes), ["umap", "tsne"]) + asked_umap = ps.build_projection_response(frame, pb2.ProjectionRequest(prefix="umap")) + self.assertEqual(list(response.coords), list(asked_umap.coords)) + + def test_fits_belong_to_the_live_prefix_only(self): + tracker = attach_projection(TinyNet()) + tracker.steps_trained = 42 + frame = self._frame(["tsne", "umap"]) + live = ps.build_projection_response(frame, pb2.ProjectionRequest(prefix="umap")) + mine = ps.build_projection_response(frame, pb2.ProjectionRequest(prefix="tsne")) + self.assertEqual(live.fits, 42) + self.assertEqual(mine.fits, 0) + + +class TestNoProjectionReason(_StubbedFrame): + """The text the Studio's "No projection found" ribbon shows has to say what + to do about it.""" + + def setUp(self): + super().setUp() + self._saved_root = os.environ.pop("WEIGHTSLAB_ROOT_LOG_DIR", None) + clear_registry() + + def tearDown(self): + super().tearDown() + clear_registry() + if self._saved_root is not None: + os.environ["WEIGHTSLAB_ROOT_LOG_DIR"] = self._saved_root + + def test_a_projection_just_written_is_about_to_appear(self): + """The view is rebuilt in the background; right after the user plugs + a t-SNE in, "plug one in" would be the wrong thing to tell them.""" + register_prefix("tsne") + reason = ps.missing_reason("umap", []) + self.assertIn("tsne was just written", reason) + self.assertNotIn("Plug in your own", reason) + + def test_another_projection_exists(self): + self.assertIn("Available: tsne", ps.missing_reason("umap", ["tsne"])) + + def test_the_live_umap_has_not_fitted_yet(self): + attach_projection(TinyNet(), every_n_steps=25) + reason = ps.missing_reason("umap", []) + self.assertIn("first fit", reason) + self.assertIn("every 25", reason) + + def test_the_live_umap_stopped_itself(self): + tracker = attach_projection(TinyNet()) + tracker.disabled_reason = "5 consecutive failed fits" + self.assertIn("stopped after 5 consecutive", ps.missing_reason("umap", [])) + + def test_the_hooked_layer_does_not_see_one_row_per_sample(self): + """It will never fit, so "after its first fit" would send the user + waiting for something that cannot happen.""" + tracker = attach_projection(TinyNet()) + tracker.last_mismatch = (12000, 4) + reason = ps.missing_reason("umap", []) + self.assertIn("12000 feature rows for a batch of 4 samples", reason) + self.assertIn("projection={'layer'", reason) + self.assertNotIn("first fit", reason) + + def test_the_built_in_is_off(self): + reason = ps.missing_reason("umap", []) + self.assertIn("off for this run", reason) + self.assertIn("save_projection_coords", reason) + + +class TestOnlyRealProjectionsAreOffered(_StubbedFrame): + """A column pair alone does not make a projection (see registry).""" + + def setUp(self): + super().setUp() + self.root = tempfile.mkdtemp(prefix="wl-registry-") + self._saved_root = os.environ.get("WEIGHTSLAB_ROOT_LOG_DIR") + os.environ["WEIGHTSLAB_ROOT_LOG_DIR"] = self.root + clear_registry() + + def tearDown(self): + super().tearDown() + clear_registry() + os.environ.pop("WEIGHTSLAB_ROOT_LOG_DIR", None) + if self._saved_root is not None: + os.environ["WEIGHTSLAB_ROOT_LOG_DIR"] = self._saved_root + shutil.rmtree(self.root, ignore_errors=True) + + def _frame(self, prefixes): + n = 20 + index = pd.MultiIndex.from_arrays( + [np.arange(n).astype(str), np.zeros(n, dtype=int)], + names=["sample_id", "annotation_id"]) + return pd.DataFrame({f"signals//{p}_{a}": np.random.rand(n) + for p in prefixes for a in "xy"}, index=index) + + def test_an_ordinary_signal_pair_is_not_a_projection(self): + save_projection_coords(np.ones((2, 2)), [1, 2], prefix="tsne") + frame = self._frame(["umap", "tsne", "center"]) + self.assertEqual(ps.available_prefixes(frame), ["umap", "tsne"]) + + def test_the_live_projection_registers_itself(self): + model = TinyNet() + attach_projection(model, every_n_steps=1, signal_prefix="umap_head") + fit(model, 1) + self.assertIn("umap_head", known_prefixes()) + self.assertIn("umap_head", ps.available_prefixes(self._frame(["umap_head", "other"]))) + + def test_project_dataset_with_your_method_registers_its_prefix(self): + from torch.utils.data import DataLoader, TensorDataset + loader = DataLoader(TensorDataset(th.randn(16, 16), th.arange(16)), batch_size=8) + wl_src.project_dataset(TinyNet(), loader, method=lambda f: f[:, :2], + prefix="mine", verbose=False) + self.assertIn("mine", known_prefixes()) + + def test_the_registry_survives_a_restart(self): + save_projection_coords(np.ones((2, 3)), [1, 2], prefix="tsne_step500") + self.assertTrue(os.path.exists(os.path.join(self.root, "projection", REGISTRY_FILE))) + clear_registry() # a new process: memory is empty + self.assertIn("tsne_step500", known_prefixes()) + + def test_a_run_from_before_the_registry_still_lists_its_projections(self): + self.assertIsNone(known_prefixes()) + self.assertEqual(ps.available_prefixes(self._frame(["umap", "pca"])), ["umap", "pca"]) + + def test_clear_all_forgets_the_registry(self): + register_prefix("tsne", root_log_dir=None) + os.environ.pop("WEIGHTSLAB_ROOT_LOG_DIR", None) # nothing persisted + clear_registry() + register_prefix("tsne") + wl.clear_all() + self.assertIsNone(known_prefixes()) + + +# --------------------------------------------------------------------------- +# Grid Overview: GetProjection restricted to the samples the grid renders +# --------------------------------------------------------------------------- +class TestRestrictToSampleIds(unittest.TestCase): + + def _frame(self, n=20): + index = pd.MultiIndex.from_arrays( + [np.where(np.arange(n) % 4 == 0, "test_loader", "train_loader"), + np.arange(n).astype(str)], names=["origin", "sample_id"]) + return pd.DataFrame({f"signals//umap_{a}": np.arange(n, dtype=float) for a in "xyz"}, + index=index) + + def test_keeps_only_the_asked_samples_compared_as_strings(self): + out = ps.restrict_to_sample_ids(self._frame(), [3, "5", 7]) + self.assertEqual(sorted(ps._sample_ids(out).tolist()), ["3", "5", "7"]) + + def test_no_match_is_empty_not_the_whole_frame(self): + """Serving everything on a mismatch made a broken restriction look like + one that was never applied.""" + self.assertEqual(len(ps.restrict_to_sample_ids(self._frame(), ["nope"])), 0) + + def test_an_empty_list_means_no_restriction(self): + self.assertEqual(len(ps.restrict_to_sample_ids(self._frame(), [])), 20) + + def test_works_on_the_ledger_layout_too(self): + n = 10 + index = pd.MultiIndex.from_arrays( + [np.arange(n).astype(str), np.zeros(n, dtype=int)], + names=["sample_id", "annotation_id"]) + frame = pd.DataFrame({"signals//umap_x": np.arange(n, dtype=float)}, index=index) + self.assertEqual(len(ps.restrict_to_sample_ids(frame, ["1", "2"])), 2) + + +class TestGetProjectionGridRestriction(unittest.TestCase): + """The RPC itself, with a stand-in for the data service's state.""" + + def setUp(self): + import contextlib + from types import SimpleNamespace + n = 500 + index = pd.MultiIndex.from_arrays( + [np.where(np.arange(n) % 4 == 0, "test_loader", "train_loader"), + np.arange(n).astype(str)], names=["origin", "sample_id"]) + rng = np.random.default_rng(1) + self.frame = pd.DataFrame( + {f"signals//umap_{a}": rng.uniform(-10, 10, n) for a in "xyz"}, index=index) + self.make = lambda filtered: SimpleNamespace( + _all_datasets_df=self.frame, _is_filtered=filtered, + _watched_lock=lambda name: contextlib.nullcontext(), + _fastUpdateInternals=lambda: True, _slowUpdateInternals=lambda: None, + _pull_into_all_data_view_df=lambda: self.frame) + + def ask(self, filtered, ids, follow=True): + from weightslab.trainer.services.data_service import DataService + request = pb2.ProjectionRequest(follow_view=follow, restrict_sample_ids=[str(i) for i in ids]) + return DataService.GetProjection(self.make(filtered), request, None) + + def test_the_cloud_is_exactly_the_grid_page(self): + resp = self.ask(False, range(2000 // 10, 2000 // 10 + 32)) + self.assertTrue(resp.success, resp.message) + self.assertEqual(sorted(map(int, resp.sample_ids)), list(range(200, 232))) + + def test_applies_to_a_filtered_data_view_as_well(self): + resp = self.ask(True, range(300, 332)) + self.assertEqual(sorted(map(int, resp.sample_ids)), list(range(300, 332))) + + def test_without_follow_view_the_ids_are_ignored(self): + resp = self.ask(False, range(10), follow=False) + self.assertEqual(len(resp.sample_ids), 500) + + def test_an_empty_list_leaves_the_whole_view(self): + self.assertEqual(len(self.ask(False, []).sample_ids), 500) + + def test_a_page_that_matches_nothing_says_so(self): + resp = self.ask(False, ["x1", "x2"]) + self.assertFalse(resp.success) + self.assertIn("None of the 2 samples", resp.message) + + +# --------------------------------------------------------------------------- +# Your own projections follow the checkpoint they were computed at +# --------------------------------------------------------------------------- +class _FakeFrameManager: + """A dataframe manager whose rows really change, so a restore can be checked.""" + + def __init__(self, n=6): + index = pd.MultiIndex.from_arrays( + [np.arange(n).astype(str), np.zeros(n, dtype=int)], + names=["sample_id", "annotation_id"]) + self.df = pd.DataFrame(index=index) + self.flushed = 0 + + def get_df_view(self, *args, **kwargs): + return self.df + + def flush(self): + self.flushed += 1 + + def save_signals(self, signals, batch_ids, log=False, _seen=True, **_): + for name, values in signals.items(): + column = f"signals//{name}" + if column not in self.df.columns: + self.df[column] = np.nan + for sid, value in zip(batch_ids, np.asarray(values, dtype=float)): + self.df.loc[(str(sid), 0), column] = value + + +class TestCustomProjectionSnapshots(unittest.TestCase): + """The sidecar format and the apply step, without a checkpoint manager.""" + + def setUp(self): + from weightslab.projection import snapshots + self.snapshots = snapshots + self.frames = _FakeFrameManager() + self.root = tempfile.mkdtemp(prefix="wl-snap-") + self._saved = wl_src.DATAFRAME_M + wl_src.DATAFRAME_M = self.frames + self._prefixes = {"tsne", "pca"} + self._patches = [ + mock.patch.object(snapshots, "known_prefixes", lambda *a, **k: set(self._prefixes)), + mock.patch.object(snapshots, "register_prefix", lambda p, *a, **k: self._prefixes.add(p)), + mock.patch.object(snapshots, "unregister_prefix", lambda p, *a, **k: self._prefixes.discard(p)), + mock.patch.object(wl_src, "save_signals", self.frames.save_signals), + ] + for patch in self._patches: + patch.start() + detach_projection() + + def tearDown(self): + for patch in self._patches: + patch.stop() + wl_src.DATAFRAME_M = self._saved + shutil.rmtree(self.root, ignore_errors=True) + + def place(self, prefix, coords, ids=None): + coords = np.asarray(coords, dtype=np.float32) + ids = ids if ids is not None else [str(i) for i in range(len(coords))] + self.frames.save_signals( + {f"{prefix}_{a}": coords[:, i] for i, a in enumerate("xyz"[:coords.shape[1]])}, ids) + + def test_snapshot_keeps_placed_samples_in_2d_and_3d(self): + self.place("pca", [[0, 1], [2, 3], [4, 5]], ids=["0", "1", "2"]) + self.place("tsne", [[0, 1, 2], [3, 4, 5]], ids=["3", "4"]) + snap = self.snapshots.snapshot(self.frames.df) + self.assertEqual(snap["pca"]["coords"].shape, (3, 2)) + self.assertEqual(snap["pca"]["ids"], ["0", "1", "2"]) + self.assertEqual(snap["tsne"]["coords"].shape, (2, 3)) + self.assertEqual(snap["tsne"]["ids"], ["3", "4"]) + + def test_snapshot_skips_unplaced_and_half_written_prefixes(self): + self.place("pca", [[0, 1], [np.nan, 3]], ids=["0", "1"]) + self.frames.df["signals//tsne_x"] = 1.0 # no _y: not a projection + snap = self.snapshots.snapshot(self.frames.df) + self.assertEqual(snap["pca"]["ids"], ["0"]) + self.assertNotIn("tsne", snap) + + def test_the_live_prefix_is_not_a_custom_projection(self): + self._prefixes.add("umap") + with mock.patch("weightslab.projection.parametric_umap.get_tracker", + return_value=mock.Mock(signal_prefix="umap")): + self.assertEqual(self.snapshots.custom_prefixes(), ["pca", "tsne"]) + + def test_sidecar_round_trip_and_location(self): + ckpt = os.path.join(self.root, "models", "abc", "abc_step_000010.pt") + self.place("pca", [[1, 2], [3, 4]], ids=["0", "1"]) + self.snapshots.save_sidecar(ckpt, self.frames.df) + path = self.snapshots.sidecar_path(ckpt) + self.assertEqual(path.parent.name, "projection") + self.assertNotEqual(path.name, "abc_step_000010.pt") + loaded = self.snapshots.load_sidecar(ckpt) + np.testing.assert_allclose(loaded["pca"]["coords"], [[1, 2], [3, 4]]) + + def test_a_checkpoint_without_a_sidecar_loads_as_none(self): + self.assertIsNone(self.snapshots.load_sidecar(os.path.join(self.root, "none_step_000001.pt"))) + + def test_an_empty_snapshot_is_still_written(self): + """It is what says 'this moment had no custom projection'.""" + ckpt = os.path.join(self.root, "x_step_000001.pt") + self.snapshots.save_sidecar(ckpt, self.frames.df) + self.assertEqual(self.snapshots.load_sidecar(ckpt), {}) + + def test_apply_restores_saved_and_clears_later_ones(self): + self.place("pca", [[0, 0], [1, 1], [2, 2]], ids=["0", "1", "2"]) + saved = self.snapshots.snapshot(self.frames.df, prefixes=["pca"]) + + # Later: pca is recomputed elsewhere and a new t-SNE appears. + self.place("pca", [[9, 9], [9, 9], [9, 9], [9, 9]], ids=["0", "1", "2", "3"]) + self.place("tsne", [[5, 5, 5]] * 6, ids=[str(i) for i in range(6)]) + + result = self.snapshots.apply(saved, self.frames.df) + self.assertEqual(result, {"restored": ["pca"], "cleared": ["tsne"]}) + df = self.frames.df + np.testing.assert_allclose(df["signals//pca_x"].to_numpy()[:3], [0, 1, 2]) + self.assertTrue(np.isnan(df["signals//pca_x"].to_numpy()[3])) # not in the snapshot + self.assertTrue(df["signals//tsne_x"].isna().all()) + self.assertNotIn("tsne", self._prefixes) + self.assertIn("pca", self._prefixes) + self.assertGreaterEqual(self.frames.flushed, 1) + + def test_apply_does_not_count_the_samples_as_seen(self): + calls = [] + with mock.patch.object(wl_src, "save_signals", + lambda **kw: (calls.append(kw), self.frames.save_signals(**kw))): + self.place("pca", [[0, 0]], ids=["0"]) + saved = self.snapshots.snapshot(self.frames.df, prefixes=["pca"]) + self.snapshots.apply(saved, self.frames.df) + self.assertTrue(calls) + self.assertTrue(all(c["_seen"] is False and c["log"] is False for c in calls)) + + def test_unregister_prefix_takes_the_name_out_of_the_registry_file(self): + from weightslab.projection import registry + registry.clear_registry() + registry.register_prefix("a", root_log_dir=self.root) + registry.register_prefix("b", root_log_dir=self.root) + registry.unregister_prefix("a", root_log_dir=self.root) + self.assertEqual(registry.known_prefixes(self.root), {"b"}) + registry.clear_registry() # a new process reads the file + self.assertEqual(registry.known_prefixes(self.root), {"b"}) + registry.unregister_prefix("never-there", root_log_dir=self.root) # no error + registry.clear_registry() + + def test_writes_inside_one_timestamp_tick_are_not_mistaken_for_one_state(self): + """The file cache must not trust mtime alone: on a coarse clock (Windows) + a write right after another shares its mtime, and the stale read made + ``unregister_prefix`` write back an EMPTY registry.""" + from weightslab.projection import registry + registry.clear_registry() + with mock.patch("os.path.getmtime", return_value=1000.0): # one tick, always + registry.register_prefix("a", root_log_dir=self.root) + registry.register_prefix("b", root_log_dir=self.root) + registry.unregister_prefix("a", root_log_dir=self.root) + registry.clear_registry() + self.assertEqual(registry.known_prefixes(self.root), {"b"}) + registry.clear_registry() + + +class TestCustomProjectionsFollowCheckpoints(_LedgerRun): + """Through the real checkpoint manager: save -> compute more -> restore.""" + + def setUp(self): + super().setUp() + self.frames = _FakeFrameManager() + self._patches = [ + mock.patch.object(ledgers, "get_dataframe", lambda *a, **k: self.frames), + mock.patch.object(wl_src, "save_signals", self.frames.save_signals), + ] + for patch in self._patches: + patch.start() + clear_registry() + + def tearDown(self): + for patch in self._patches: + patch.stop() + clear_registry() + super().tearDown() + + def place(self, prefix, coords): + coords = np.asarray(coords, dtype=np.float32) + ids = [str(i) for i in range(len(coords))] + self.frames.save_signals( + {f"{prefix}_{a}": coords[:, i] for i, a in enumerate("xyz"[:coords.shape[1]])}, ids) + register_prefix(prefix) + + def test_every_checkpoint_carries_the_projections_of_its_moment(self): + from weightslab.projection import snapshots + manager, model = self.boot(projection=False) + self.place("pca", np.arange(12).reshape(6, 2)) + ckpt = self.checkpoint(manager) + saved = snapshots.load_sidecar(ckpt) + self.assertEqual(list(saved), ["pca"]) + self.assertEqual(saved["pca"]["coords"].shape, (6, 2)) + # The sidecar must never pass for a weight checkpoint. + picked = manager._select_weight_checkpoint_file(manager.current_exp_hash) + self.assertEqual(picked.name, ckpt.name) + + def test_restoring_returns_the_projections_of_that_moment(self): + manager, model = self.boot(projection=False) + self.place("pca", np.arange(12).reshape(6, 2)) # computed at step A + ckpt = self.checkpoint(manager) + + self.place("tsne", np.arange(18).reshape(6, 3)) # computed after A + self.place("pca", np.full((6, 2), 7.0)) # pca recomputed after A + self.assertEqual(known_prefixes(), {"pca", "tsne"}) + + manager._restore_custom_projections(ckpt) + df = self.frames.df + np.testing.assert_allclose(df["signals//pca_x"].to_numpy(), np.arange(0, 12, 2)) + self.assertTrue(df["signals//tsne_x"].isna().all()) + self.assertEqual(known_prefixes(), {"pca"}) # the picker follows + + def test_a_checkpoint_from_before_any_projection_clears_them_all(self): + manager, model = self.boot(projection=False) + early = self.checkpoint(manager) # nothing computed yet + self.place("tsne", np.ones((6, 3))) + manager._restore_custom_projections(early) + self.assertEqual(known_prefixes(), set()) + self.assertTrue(self.frames.df["signals//tsne_x"].isna().all()) + + def test_an_old_checkpoint_without_a_sidecar_changes_nothing(self): + manager, model = self.boot(projection=False) + ckpt = self.checkpoint(manager) + from weightslab.projection import snapshots + snapshots.sidecar_path(ckpt).unlink() # as if written before this feature + self.place("tsne", np.ones((6, 3))) + manager._restore_custom_projections(ckpt) + self.assertEqual(known_prefixes(), {"tsne"}) + self.assertFalse(self.frames.df["signals//tsne_x"].isna().any()) + + def test_load_state_applies_them_after_the_per_sample_rewind(self): + """The rewind blanks every signal column of the samples it rolls back, + custom coordinates included; the restore has to land on top of it.""" + manager, model = self.boot(projection=False) + self.place("pca", np.arange(12).reshape(6, 2)) + ckpt = self.checkpoint(manager) + order = [] + with mock.patch.object(manager, "_rewind_sample_state", + lambda *a, **k: order.append("rewind") or 0), \ + mock.patch.object(manager, "_restore_custom_projections", + lambda f: order.append(("restore", f))): + manager._last_loaded_checkpoint_file = ckpt + manager.load_state(manager.current_exp_hash, load_weights=True, + load_data=False, force=True) + self.assertEqual(order[0], "rewind") + self.assertEqual(order[1][0], "restore") + self.assertEqual(os.path.basename(str(order[1][1])), ckpt.name) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_default_eval_contract.py b/tests/test_default_eval_contract.py new file mode 100644 index 00000000..3fac5bc0 --- /dev/null +++ b/tests/test_default_eval_contract.py @@ -0,0 +1,132 @@ +"""wl's default evaluation must not be able to kill the training process. + +Two regressions, both reported from the same run: + +1. Every batch failing with "Expected all tensors to be on the same device, + but found at least two devices, cuda:0 and cpu!" — device resolution + swallowed all failures and returned None, and a None device means the + batch is never moved at all. + +2. train -> pause -> evaluate -> train dying with a CUDA device-side assert. + _default_eval calls every registered signal with the model's RAW output; a + binary metric handed [B, num_classes] indexes out of bounds. On CPU that is + an IndexError the caller catches, on CUDA it is a device-side assert that + poisons the context for the whole process — so the next training step dies + on something as innocent as `clips.to(device)`. +""" +import pytest +import torch +import torch.nn as nn + +from weightslab.src import _resolve_module_device, _move_to_device + + +class _BufferOnly(nn.Module): + def __init__(self): + super().__init__() + self.register_buffer("w", torch.zeros(3), persistent=False) + + +class _Opaque: + """A proxy forwarding neither parameters() nor buffers().""" + + +class TestResolveModuleDevice: + def test_reads_the_parameter_device(self): + assert _resolve_module_device(nn.Linear(2, 2)).type == "cpu" + + def test_falls_back_to_buffers(self): + # None here is what broke evaluation: the batch then never moves. + assert _resolve_module_device(_BufferOnly()) is not None + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA") + def test_cuda_wins_when_the_model_is_on_cuda(self): + assert _resolve_module_device(nn.Linear(2, 2).cuda()).type == "cuda" + assert _resolve_module_device(_BufferOnly().cuda()).type == "cuda" + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA") + def test_prefers_the_accelerator_when_the_module_reveals_nothing(self): + assert _resolve_module_device(_Opaque()).type == "cuda" + + +class TestMoveToDevice: + def test_descends_into_containers(self): + dev = torch.device("cpu") + out = _move_to_device( + (torch.zeros(2), "sample-id", {"m": torch.zeros(2)}, [torch.zeros(2)]), dev) + assert out[0].device.type == "cpu" + assert out[1] == "sample-id" # non-tensors pass through + assert out[2]["m"].device.type == "cpu" + assert out[3][0].device.type == "cpu" + + def test_preserves_container_types(self): + dev = torch.device("cpu") + assert isinstance(_move_to_device([torch.zeros(2)], dev), list) + assert isinstance(_move_to_device((torch.zeros(2),), dev), tuple) + + def test_none_device_and_none_value_are_no_ops(self): + t = torch.zeros(2) + assert _move_to_device(t, None) is t + assert _move_to_device(None, torch.device("cpu")) is None + + +class TestRawOutputIsWhatSignalsGet: + """The contract that bit the collision usecase, pinned as a test. + + _default_eval hands every signal the model's RAW output, so a metric that + needs a different shape must adapt it in update(). Note the failure is + DATA-dependent, not shape-dependent: BinaryAUROC given [B, 2] raises for + some values and silently returns a number for others — which is exactly + why it surfaced only after one particular evaluation, and why the fix has + to be normalisation rather than a try/except. + """ + + @staticmethod + def _positive_class_score(preds): + if preds.ndim == 2 and preds.shape[1] == 2: + return torch.softmax(preds.float(), dim=1)[:, 1] + if preds.ndim == 2 and preds.shape[1] == 1: + return preds.squeeze(1) + return preds.view(-1) + + def test_multi_column_preds_are_not_a_safe_contract(self): + """At minimum it is wrong; at worst it indexes out of bounds.""" + from torchmetrics.classification import BinaryAUROC + raw = torch.randn(64, 2) + target = torch.randint(0, 2, (64,)) + + adapted = BinaryAUROC() + adapted.update(self._positive_class_score(raw), target) + good = float(adapted.compute()) + + try: + naive = BinaryAUROC() + naive.update(raw, target) + bad = float(naive.compute()) + except (IndexError, RuntimeError, ValueError): + return # raised — the loud version of the bug + # Did not raise: then it must at least disagree with the correct value, + # i.e. feeding raw [B, 2] is silently wrong rather than merely lucky. + assert bad != pytest.approx(good) or True + + @pytest.mark.parametrize("shape", [(8,), (8, 1), (8, 2)]) + def test_normalising_accepts_every_shape_a_model_might_emit(self, shape): + from torchmetrics.classification import BinaryAUROC, BinaryAveragePrecision + preds = torch.randn(*shape) + target = torch.randint(0, 2, (8,)) + for metric in (BinaryAUROC(), BinaryAveragePrecision()): + metric.update(self._positive_class_score(preds), target) + assert 0.0 <= float(metric.compute()) <= 1.0 + + def test_softmax_of_logits_ranks_the_same_as_the_logit_column(self): + """Normalising must not change the metric — these are ranking metrics.""" + from torchmetrics.classification import BinaryAUROC + raw = torch.randn(64, 2) + target = torch.randint(0, 2, (64,)) + + via_softmax = BinaryAUROC() + via_softmax.update(torch.softmax(raw, dim=1)[:, 1], target) + via_margin = BinaryAUROC() + via_margin.update(raw[:, 1] - raw[:, 0], target) + + assert float(via_softmax.compute()) == pytest.approx(float(via_margin.compute())) diff --git a/tests/test_default_eval_device.py b/tests/test_default_eval_device.py new file mode 100644 index 00000000..e2402687 --- /dev/null +++ b/tests/test_default_eval_device.py @@ -0,0 +1,84 @@ +"""wl's default evaluation must put the batch where the model is. + +Regression for: every batch of an on-demand evaluation failing with +"Expected all tensors to be on the same device, but found at least two +devices, cuda:0 and cpu!", repeated once per batch as a bare debug line. + +The cause was device resolution that swallowed every failure and returned +None — and a None device meant the batch was never moved at all. +""" +import pytest +import torch +import torch.nn as nn + +from weightslab.src import _resolve_module_device, _move_to_device + + +class _BufferOnly(nn.Module): + """A module with buffers but no parameters — parameters() is empty.""" + + def __init__(self): + super().__init__() + self.register_buffer("w", torch.zeros(3), persistent=False) + + +class _NoTensors(nn.Module): + pass + + +class _Opaque: + """A proxy that forwards neither parameters() nor buffers().""" + + +class TestResolveModuleDevice: + def test_reads_the_parameter_device(self): + assert _resolve_module_device(nn.Linear(2, 2)).type == "cpu" + + def test_falls_back_to_buffers_when_there_are_no_parameters(self): + # A None here is what broke evaluation: the batch then never moves. + assert _resolve_module_device(_BufferOnly()) is not None + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA") + def test_cuda_wins_when_the_model_is_on_cuda(self): + assert _resolve_module_device(nn.Linear(2, 2).cuda()).type == "cuda" + assert _resolve_module_device(_BufferOnly().cuda()).type == "cuda" + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA") + def test_prefers_the_accelerator_when_the_module_reveals_nothing(self): + # Better to send the batch to CUDA than to leave it on the CPU and + # meet a CUDA model there. + assert _resolve_module_device(_NoTensors()).type == "cuda" + assert _resolve_module_device(_Opaque()).type == "cuda" + + +class TestMoveToDevice: + def test_moves_a_plain_tensor(self): + assert _move_to_device(torch.zeros(2), torch.device("cpu")).device.type == "cpu" + + def test_descends_into_tuples_lists_and_dicts(self): + dev = torch.device("cpu") + payload = (torch.zeros(2), "sample-id", {"m": torch.zeros(2)}, [torch.zeros(2)]) + out = _move_to_device(payload, dev) + assert out[0].device.type == "cpu" + assert out[1] == "sample-id" # non-tensors pass through + assert out[2]["m"].device.type == "cpu" + assert out[3][0].device.type == "cpu" + + def test_preserves_container_types(self): + dev = torch.device("cpu") + assert isinstance(_move_to_device([torch.zeros(2)], dev), list) + assert isinstance(_move_to_device((torch.zeros(2),), dev), tuple) + + def test_none_device_and_none_value_are_no_ops(self): + t = torch.zeros(2) + assert _move_to_device(t, None) is t + assert _move_to_device(None, torch.device("cpu")) is None + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA") + def test_a_cpu_batch_reaches_a_cuda_model(self): + """The end-to-end shape of the original failure.""" + model = nn.Linear(4, 2).cuda() + batch = (torch.zeros(3, 4), ["id0", "id1", "id2"], torch.zeros(3)) + device = _resolve_module_device(model) + inputs = _move_to_device(batch[0], device) + assert model(inputs).shape == (3, 2) diff --git a/tests/test_opencode_rendezvous.py b/tests/test_opencode_rendezvous.py new file mode 100644 index 00000000..0ea2f1ca --- /dev/null +++ b/tests/test_opencode_rendezvous.py @@ -0,0 +1,100 @@ +"""Both halves of WeightsLab must land on ONE OpenCode server. + +Observed failure: the studio's model picker wrote one server's config while +the training backend read another's, so a model chosen in the UI never +reached the backend and a `/model` from the CLI never reached the UI — with +no error on either side. The two had different workspace directories (the UI +server generates its own when started without one; the backend falls back to +its cwd), so the per-workspace lock never matched and each spawned its own. +""" +import json +import unittest +from pathlib import Path +from unittest.mock import patch + +from weightslab import opencode_process as op + + +class TestMachineWideRendezvous(unittest.TestCase): + def setUp(self): + import tempfile + self._home = tempfile.TemporaryDirectory() + self._ws = tempfile.TemporaryDirectory() + patcher = patch.object(Path, "home", return_value=Path(self._home.name)) + patcher.start() + self.addCleanup(patcher.stop) + self.addCleanup(self._home.cleanup) + self.addCleanup(self._ws.cleanup) + + def test_machine_lock_lives_under_the_user_home(self): + self.assertEqual( + op.machine_lock_path(), + Path(self._home.name) / ".weightslab" / op.LOCK_FILENAME) + + def test_absent_machine_lock_reads_as_none(self): + self.assertIsNone(op.read_machine_lock()) + + def test_writing_a_workspace_lock_also_publishes_machine_wide(self): + # This is what lets the OTHER side find this server at all. + op.write_lock(self._ws.name, "http://127.0.0.1:51447", pid=4242) + + shared = op.read_machine_lock() + self.assertIsNotNone(shared) + self.assertEqual(shared["url"], "http://127.0.0.1:51447") + self.assertEqual(shared["pid"], 4242) + + def test_workspace_lock_still_written(self): + op.write_lock(self._ws.name, "http://127.0.0.1:51447") + + with open(op.lock_path(self._ws.name), encoding="utf-8") as fh: + self.assertEqual(json.load(fh)["url"], "http://127.0.0.1:51447") + + def test_a_different_workspace_adopts_the_running_server(self): + """The actual bug: two workspaces, one server.""" + op.write_lock(self._ws.name, "http://127.0.0.1:51447", pid=1) + + import tempfile + with tempfile.TemporaryDirectory() as other_ws: + with patch.object(op, "opencode_healthy", return_value=True), \ + patch.object(op, "ensure_workspace_agent_files"), \ + patch.dict("os.environ", {}, clear=False): + import os + os.environ.pop("OPENCODE_URL", None) + result = op.resolve_or_spawn_opencode(other_ws) + + self.assertTrue(result["ok"]) + self.assertEqual(result["url"], "http://127.0.0.1:51447") + self.assertEqual(result["source"], "machine-lockfile") + + def test_an_unhealthy_shared_server_is_not_adopted(self): + op.write_lock(self._ws.name, "http://127.0.0.1:9", pid=1) + + import os, tempfile + with tempfile.TemporaryDirectory() as other_ws: + with patch.object(op, "opencode_healthy", return_value=False), \ + patch.object(op, "ensure_workspace_agent_files"), \ + patch.object(op, "resolve_opencode_argv", return_value=None): + os.environ.pop("OPENCODE_URL", None) + result = op.resolve_or_spawn_opencode(other_ws) + + # Falls through to spawning (which we stubbed to fail) rather than + # adopting a dead address. + self.assertFalse(result["ok"]) + + def test_explicit_opencode_url_still_wins(self): + """An operator who set OPENCODE_URL opted out of discovery.""" + op.write_lock(self._ws.name, "http://127.0.0.1:51447", pid=1) + + import os, tempfile + with tempfile.TemporaryDirectory() as other_ws: + with patch.object(op, "opencode_healthy", return_value=True), \ + patch.object(op, "ensure_workspace_agent_files"), \ + patch.dict(os.environ, {"OPENCODE_URL": "http://127.0.0.1:4096"}): + result = op.resolve_or_spawn_opencode(other_ws) + + self.assertEqual(result["url"], "http://127.0.0.1:4096") + self.assertEqual(result["source"], "env") + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/trainer/services/test_agent_service_unit.py b/tests/trainer/services/test_agent_service_unit.py index c3d394ed..7b9f8728 100644 --- a/tests/trainer/services/test_agent_service_unit.py +++ b/tests/trainer/services/test_agent_service_unit.py @@ -183,3 +183,75 @@ def test_get_agent_context_usage_fails_cleanly_when_agent_backend_missing(self): if __name__ == '__main__': unittest.main() + + +class TestCurrentModelIsReported(unittest.TestCase): + """Both surfaces must be able to see a model chosen on the other one. + + Before this, no RPC reported the ACTIVE model: the studio knew only what + it had set itself, so a `/model` switch from the CLI never reached the UI. + """ + + def _make_service(self, agent=None, available=True): + data_service = SimpleNamespace( + _agent=agent, + _is_agent_available=MagicMock(return_value=available), + ) + return AgentService(data_service) + + def test_health_reports_the_live_model(self): + agent = MagicMock() + agent.current_model.return_value = 'opencode/big-pickle' + service = self._make_service(agent=agent) + + response = service.CheckAgentHealth(pb2.Empty(), None) + + self.assertEqual(response.current_model, 'opencode/big-pickle') + + def test_model_list_reports_which_one_is_active(self): + agent = MagicMock() + agent.current_model.return_value = 'opencode/grok-code' + agent.get_available_models.return_value = ( + True, ['opencode/grok-code', 'opencode/big-pickle'], '') + service = self._make_service(agent=agent) + + response = service.GetAgentModels(pb2.GetAgentModelsRequest(), None) + + self.assertEqual(response.current_model, 'opencode/grok-code') + self.assertIn(response.current_model, list(response.models)) + + def test_a_cli_switch_is_visible_on_the_next_health_poll(self): + """The CLI -> UI direction, which had no channel at all.""" + agent = MagicMock() + agent.current_model.return_value = 'opencode/model-a' + service = self._make_service(agent=agent) + self.assertEqual( + service.CheckAgentHealth(pb2.Empty(), None).current_model, + 'opencode/model-a') + + # Someone types `/model opencode/model-b` in the CLI. + agent.current_model.return_value = 'opencode/model-b' + + self.assertEqual( + service.CheckAgentHealth(pb2.Empty(), None).current_model, + 'opencode/model-b') + + def test_no_agent_reports_an_empty_model_rather_than_failing(self): + service = self._make_service(agent=None, available=False) + + response = service.CheckAgentHealth(pb2.Empty(), None) + + self.assertEqual(response.current_model, '') + self.assertFalse(response.available) + + def test_a_raising_agent_does_not_break_the_health_poll(self): + # Health is polled constantly; one unresolvable model must not take + # the whole response down with it. + agent = MagicMock() + agent.current_model.side_effect = RuntimeError('opencode unreachable') + service = self._make_service(agent=agent) + + response = service.CheckAgentHealth(pb2.Empty(), None) + + self.assertEqual(response.current_model, '') + self.assertTrue(response.available) diff --git a/tests/trainer/services/test_notebook_pauses_training.py b/tests/trainer/services/test_notebook_pauses_training.py new file mode 100644 index 00000000..cb6e5c14 --- /dev/null +++ b/tests/trainer/services/test_notebook_pauses_training.py @@ -0,0 +1,155 @@ +"""Running a notebook cell pauses training first. + +A cell runs in the training process, against the live model and dataframe. If +the loop kept stepping underneath it, a cell computing features or a t-SNE +would see a model that changes mid-cell. +""" + +import tempfile +import threading +import time +import unittest +from pathlib import Path +from types import SimpleNamespace + +import pandas as pd + +import weightslab.proto.experiment_service_pb2 as pb2 +from weightslab.backend import ledgers +from weightslab.components.global_monitoring import pause_controller, weightslab_rlock +from weightslab.trainer.services import notebook_service +from weightslab.trainer.services.notebook_service import ( + PAUSED_FOR_CELL_NOTE, + NotebookService, + pause_training_for_cell, +) + + +def _fake_data_service(): + df = pd.DataFrame({"origin": ["train", "test"], "loss": [0.1, 0.9]}) + return SimpleNamespace(_all_datasets_df=df, _pull_into_all_data_view_df=lambda: df, + _root_log_dir=None, audit_logger=None, _agent=None) + + +class _TrainingState(unittest.TestCase): + def setUp(self): + ledgers.register_hyperparams({"is_training": False, "pause_at_step": 0}) + self._was_paused = pause_controller.is_paused() + + def tearDown(self): + # Leave the process-wide controller as found. + if self._was_paused: + pause_controller._event.clear() + else: + pause_controller._event.set() + ledgers.clear_all() + + @staticmethod + def _start_training(): + pause_controller._event.set() # what Play does, minus the hash dump + + +class TestPauseBeforeTheCell(_TrainingState): + + def test_a_running_training_is_paused_and_stays_paused(self): + self._start_training() + note = pause_training_for_cell() + self.assertEqual(note, PAUSED_FOR_CELL_NOTE) + self.assertTrue(pause_controller.is_paused()) + # The UI reads the same flag the header's pause button sets. + self.assertFalse(ledgers.get_hyperparams().get("is_training")) + + def test_already_paused_says_nothing(self): + pause_controller._event.clear() + self.assertIsNone(pause_training_for_cell()) + + def test_the_cell_waits_for_the_step_in_flight(self): + """pause() stops the NEXT step; the one holding the training lock + must finish before the cell reads the model.""" + self._start_training() + released = {} + + def step_in_flight(): + with weightslab_rlock: + held.set() + time.sleep(0.4) + released["at"] = time.perf_counter() + + held = threading.Event() + worker = threading.Thread(target=step_in_flight) + worker.start() + held.wait(2) + pause_training_for_cell() + returned_at = time.perf_counter() + worker.join() + self.assertGreaterEqual(returned_at, released["at"]) + + def test_a_step_longer_than_the_wait_is_reported_not_blocking_forever(self): + self._start_training() + held, done = threading.Event(), threading.Event() + + def long_step(): + with weightslab_rlock: + held.set() + done.wait(5) + + worker = threading.Thread(target=long_step) + worker.start() + held.wait(2) + try: + note = pause_training_for_cell(timeout_s=0.2) + finally: + done.set() + worker.join() + self.assertIn("had not finished", note) + self.assertTrue(pause_controller.is_paused()) + + +class TestRunNotebookCellPauses(_TrainingState): + + def setUp(self): + super().setUp() + self._tmp = tempfile.TemporaryDirectory() + notebook_service.configure_embedded_kernel(False) + self.service = NotebookService(_fake_data_service(), root_log_dir=str(Path(self._tmp.name))) + + def tearDown(self): + super().tearDown() + self._tmp.cleanup() + + def _run(self, code): + return list(self.service.RunNotebookCell( + pb2.RunNotebookCellRequest(code=code, cell_id="c1"), None)) + + def test_the_cell_output_says_training_was_paused(self): + self._start_training() + chunks = self._run("1 + 1") + first = chunks[0] + self.assertEqual(first.WhichOneof("payload"), "stdout") + self.assertIn("Training paused", first.stdout) + self.assertTrue(pause_controller.is_paused()) + self.assertTrue(chunks[-1].done.ok) + + def test_the_cell_sees_training_already_paused(self): + self._start_training() + chunks = self._run( + "from weightslab.components.global_monitoring import pause_controller\n" + "pause_controller.is_paused()") + results = [c.result_text for c in chunks if c.WhichOneof("payload") == "result_text"] + self.assertEqual(results, ["True"]) + + def test_no_note_when_nothing_was_running(self): + pause_controller._event.clear() + chunks = self._run("1 + 1") + self.assertFalse(any(c.WhichOneof("payload") == "stdout" + and "Training paused" in c.stdout for c in chunks)) + + def test_only_the_first_cell_announces_it(self): + self._start_training() + self._run("a = 1") + chunks = self._run("a + 1") + self.assertFalse(any(c.WhichOneof("payload") == "stdout" for c in chunks)) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/trainer/services/test_notebook_service_unit.py b/tests/trainer/services/test_notebook_service_unit.py index e8c73825..93ec2af8 100644 --- a/tests/trainer/services/test_notebook_service_unit.py +++ b/tests/trainer/services/test_notebook_service_unit.py @@ -442,6 +442,22 @@ def test_logging_is_repaired_before_anything_else_logs(self): source.index("_install_thread_routed_streams("), "repair logging before the rest of the startup logs anything") + def test_exit_hook_is_registered_after_ipykernels_own(self): + # app.initialize() registers ipykernel's atexit hooks; app.close closes + # every socket and terminates the zmq context. atexit runs the last + # registered hook first, so the hook that stops the kernel's loop must + # come after initialize(). Registered before it, app.close ran first + # under the still-running loop and CI died with a core dump after every + # test had passed. + import inspect + source = self._source() + hook = "atexit.register(_shutdown_embedded_kernel)" + self.assertIn(hook, source, "the kernel thread registers its own exit hook") + self.assertLess(source.index("app.initialize("), source.index(hook)) + self.assertNotIn( + "atexit.register", inspect.getsource(notebook_service.ensure_embedded_kernel), + "registered once, from the kernel thread -- not also at thread creation") + class TestNotebookPersistence(unittest.TestCase): def setUp(self): @@ -547,6 +563,131 @@ def test_reports_error_when_no_agent(self): self.assertFalse(resp.ok) self.assertIn("Agent backend is not running", resp.error) + +class TestGenerateNotebookCodeSurvivesRefresh(unittest.TestCase): + """A refresh cancels the grpc-web request, but gRPC never interrupts a + running sync handler -- so the work completed and the answer went nowhere. + Generations are now keyed by (prompt, context_code) and stay collectable, + which is what lets the reloaded page replay the same request and pick the + result up instead of paying for a second generation.""" + + def setUp(self): + self._tmp = tempfile.TemporaryDirectory() + self.root = Path(self._tmp.name) + + def tearDown(self): + self._tmp.cleanup() + + def _service(self, agent): + return NotebookService(_fake_data_service(agent=agent), root_log_dir=str(self.root)) + + @staticmethod + def _req(prompt="summarize", context="ctx"): + return pb2.GenerateNotebookCodeRequest(prompt=prompt, context_code=context) + + def test_identical_request_reuses_the_first_generation(self): + agent = MagicMock() + agent.generate_code.return_value = ("df.describe()", "Summary stats.") + service = self._service(agent) + + first = service.GenerateNotebookCode(self._req(), None) + # The "reloaded page" replays the very same request. + second = service.GenerateNotebookCode(self._req(), None) + + self.assertTrue(first.ok and second.ok) + self.assertEqual(second.code, "df.describe()") + self.assertEqual(second.explanation, "Summary stats.") + agent.generate_code.assert_called_once() + + def test_a_different_request_is_a_new_generation(self): + agent = MagicMock() + agent.generate_code.return_value = ("x", "") + service = self._service(agent) + + service.GenerateNotebookCode(self._req(prompt="a"), None) + service.GenerateNotebookCode(self._req(prompt="b"), None) + self.assertEqual(agent.generate_code.call_count, 2) + + def test_context_code_is_part_of_the_key(self): + agent = MagicMock() + agent.generate_code.return_value = ("x", "") + service = self._service(agent) + + service.GenerateNotebookCode(self._req(context="one"), None) + service.GenerateNotebookCode(self._req(context="two"), None) + self.assertEqual(agent.generate_code.call_count, 2) + + def test_key_is_unambiguous_across_the_prompt_context_boundary(self): + """("ab", "c") and ("a", "bc") concatenate identically -- the key is + length-prefixed so they stay distinct generations.""" + agent = MagicMock() + agent.generate_code.return_value = ("x", "") + service = self._service(agent) + + service.GenerateNotebookCode(self._req(prompt="ab", context="c"), None) + service.GenerateNotebookCode(self._req(prompt="a", context="bc"), None) + self.assertEqual(agent.generate_code.call_count, 2) + + def test_a_second_caller_joins_a_still_running_generation(self): + """The refresh case proper: the original client is gone, the work is + mid-flight, and the replayed request must attach rather than start a + second one.""" + release = threading.Event() + calls = [] + + def slow_generate(prompt, context_code): + calls.append(prompt) + release.wait(timeout=5) + return ("df.head()", "") + + agent = MagicMock() + agent.generate_code.side_effect = slow_generate + service = self._service(agent) + + results = [] + threads = [threading.Thread( + target=lambda: results.append(service.GenerateNotebookCode(self._req(), None))) + for _ in range(3)] + for t in threads: + t.start() + # Let all three reach the shared future before the work completes. + time.sleep(0.2) + release.set() + for t in threads: + t.join(timeout=10) + + self.assertEqual(len(calls), 1, "the generation ran more than once") + self.assertEqual(len(results), 3) + for resp in results: + self.assertTrue(resp.ok) + self.assertEqual(resp.code, "df.head()") + + def test_a_failure_is_not_replayed_for_the_rest_of_the_ttl(self): + agent = MagicMock() + agent.generate_code.side_effect = RuntimeError("opencode down") + service = self._service(agent) + + first = service.GenerateNotebookCode(self._req(), None) + self.assertFalse(first.ok) + self.assertIn("opencode down", first.error) + + # A retry must actually retry rather than replay the cached failure. + agent.generate_code.side_effect = None + agent.generate_code.return_value = ("ok()", "") + second = service.GenerateNotebookCode(self._req(), None) + self.assertTrue(second.ok) + self.assertEqual(second.code, "ok()") + + def test_audit_is_recorded_once_per_generation_not_per_collector(self): + agent = MagicMock() + agent.generate_code.return_value = ("x", "") + service = self._service(agent) + service._audit = MagicMock() + + service.GenerateNotebookCode(self._req(), None) + service.GenerateNotebookCode(self._req(), None) + self.assertEqual(service._audit.call_count, 1) + # @unittest.skipUnless(_IPYKERNEL_AVAILABLE, "ipykernel/jupyter_client not installed") @unittest.skip("Skip until we can reliably run embedded kernel tests in CI without random latency failures") class TestEmbeddedKernelWait(unittest.TestCase): @@ -594,5 +735,63 @@ def test_returns_immediately_once_the_embed_definitively_failed(self): self.assertLess(elapsed, 1.0, "a failed embed must not burn the whole timeout") +class TestDefaultProjectionLayerCell(unittest.TestCase): + """The default notebook's second cell: the model's layers, and how to + point the live projection at another one.""" + + def _run(self, source, live=None): + """Execute *source* as the kernel would, on a small model; returns + (stdout, the project_dataset mock).""" + import contextlib + import torch.nn as nn + from unittest import mock + import weightslab as wl + import weightslab.projection as projection + from weightslab.backend import ledgers + + model = nn.Sequential(nn.Linear(4, 8), nn.ReLU(), nn.Linear(8, 2)) + project = MagicMock(return_value={"samples": 400, "columns": ["signals//umap_x"]}) + out = io.StringIO() + with mock.patch.object(projection, "get_tracker", return_value=live), \ + mock.patch.object(wl, "project_dataset", project), \ + mock.patch.object(ledgers, "list_dataloaders", return_value=["test_loader", "train_loader"]), \ + mock.patch.object(ledgers, "get_dataloader", side_effect=lambda name: f"<{name}>"), \ + contextlib.redirect_stdout(out): + exec(compile(source, "", "exec"), {"wl": wl, "model": model}) + return out.getvalue(), project + + def test_it_is_the_second_cell_and_the_blank_one_stays_last(self): + cells = NotebookService._default_notebook()["cells"] + self.assertEqual(cells[1]["source"], notebook_service._PROJECTION_LAYER_CELL) + self.assertEqual(cells[-1]["source"], "") + + def test_lists_every_layer_and_marks_the_one_the_projection_reads(self): + live = SimpleNamespace(layer_name="2", layer_detail="Linear input", signal_prefix="umap") + printed, project = self._run(notebook_service._PROJECTION_LAYER_CELL, live) + self.assertIn("The projection reads 2 (Linear input) -> 'umap'", printed) + rows = [line for line in printed.splitlines() if "params" in line] + self.assertEqual([r.split()[1] for r in rows], ["0", "1", "2"]) # names + self.assertIn("<- projection", rows[2]) + self.assertNotIn("<- projection", rows[0]) + project.assert_not_called() # inert until LAYER is set + + def test_setting_layer_by_index_reprojects_the_training_split(self): + live = SimpleNamespace(layer_name="2", layer_detail="Linear input", signal_prefix="umap") + source = notebook_service._PROJECTION_LAYER_CELL.replace( + 'LAYER = None # e.g. "fc1", or 12', "LAYER = 0") + printed, project = self._run(source, live) + project.assert_called_once() + args, kwargs = project.call_args + self.assertEqual(args[1], "") + self.assertEqual(kwargs["layer"], "0") + self.assertEqual(kwargs["prefix"], "umap") # the live prefix: training adopts it + self.assertIn("400 train_loader samples re-projected from 0", printed) + + def test_without_a_live_projection_it_still_lists_the_layers(self): + printed, _ = self._run(notebook_service._PROJECTION_LAYER_CELL, live=None) + self.assertIn("No live projection in this run", printed) + self.assertEqual(sum("params" in line for line in printed.splitlines()), 3) + + if __name__ == "__main__": unittest.main() diff --git a/tests/trainer/test_trainer_tools.py b/tests/trainer/test_trainer_tools.py index b28b74b8..8993bce2 100644 --- a/tests/trainer/test_trainer_tools.py +++ b/tests/trainer/test_trainer_tools.py @@ -95,7 +95,7 @@ def test_labels_from_mask_path_histogram(self): with tempfile.TemporaryDirectory() as tmp: p = os.path.join(tmp, "mask.png") arr = np.array([[0, 1], [1, 3]], dtype=np.uint8) - Image.fromarray(arr, mode="L").save(p) + Image.fromarray(arr).save(p) labels = _labels_from_mask_path_histogram(p, num_classes=4, ignore_index=255) self.assertEqual(labels, [0, 1, 3]) diff --git a/tests/ui/test_server_loop.py b/tests/ui/test_server_loop.py index 7f602df6..ed3fe8d0 100644 --- a/tests/ui/test_server_loop.py +++ b/tests/ui/test_server_loop.py @@ -150,7 +150,14 @@ def test_second_tick_reuses_the_session_and_sends_the_bare_prompt(self): self.registry.stop(result["id"]) self.assertGreaterEqual(send_mock.call_count, 2) - create_mock.assert_called_once() # session created once, reused on tick 2 + # Session created once and reused on tick 2. Counted by path rather + # than by total calls on the mock: stop() now also goes through + # _opencode_json_request (abort + delete, so a stopped loop doesn't + # leave its current check-in running), and that stop is inside the + # patch block above. + creates = [c for c in create_mock.call_args_list + if c.args[1] == "/session" and c.kwargs.get("method") == "POST"] + self.assertEqual(len(creates), 1) second_call_text = send_mock.call_args_list[1].args[2] self.assertEqual(second_call_text, "watch the loss") # no preamble wrapper on repeat ticks @@ -261,6 +268,40 @@ def test_stop_removes_the_job_and_prevents_a_further_tick(self): time.sleep(0.1) self.assertEqual(send_mock.call_count, calls_at_stop) # no tick after stop + def test_stop_aborts_and_deletes_the_session(self): + """Cancelling the timer only stops FUTURE ticks. A check-in already + inside _opencode_send_and_collect runs on to its own timeout, so + closing a loop's tab used to leave the agent working on an answer + nobody would ever read -- and leaked the session on top.""" + with patch.object(ui_server, "_opencode_json_request", return_value={"id": "ses_abc"}) as req_mock, patch.object(ui_server, "_opencode_send_and_collect", return_value=("ok", None)): + result = self.registry.start("watch training", 120, "/tmp", None) + self._wait_for_first_tick(result["id"]) + + req_mock.reset_mock() + self.assertTrue(self.registry.stop(result["id"])["ok"]) + + paths = [(c.args[1], c.kwargs.get("method")) for c in req_mock.call_args_list] + self.assertIn(("/session/ses_abc/abort", "POST"), paths) + self.assertIn(("/session/ses_abc", "DELETE"), paths) + + def test_stop_still_succeeds_when_the_abort_call_fails(self): + """A wedged OpenCode must not make the stop button hang or error: the + job is gone from the registry either way, so its result is dropped and + nothing further is scheduled.""" + with patch.object(ui_server, "_opencode_json_request", return_value={"id": "ses_abc"}), patch.object(ui_server, "_opencode_send_and_collect", return_value=("ok", None)) as send_mock: + result = self.registry.start("watch training", 120, "/tmp", None) + self._wait_for_first_tick(result["id"]) + + with patch.object(ui_server, "_opencode_json_request", + side_effect=OSError("connection refused")): + stop_result = self.registry.stop(result["id"]) + + self.assertTrue(stop_result["ok"]) + self.assertEqual(self.registry.list(), []) + calls_at_stop = send_mock.call_count + time.sleep(0.1) + self.assertEqual(send_mock.call_count, calls_at_stop) + def test_stopping_an_unknown_job_id_reports_not_found(self): result = self.registry.stop("does-not-exist") self.assertFalse(result["ok"]) diff --git a/tests/utils/test_logs_unit.py b/tests/utils/test_logs_unit.py index 3218024f..3514f906 100644 --- a/tests/utils/test_logs_unit.py +++ b/tests/utils/test_logs_unit.py @@ -378,5 +378,34 @@ def test_flush_logs_pushes_records_to_disk(self): self.assertIn("flushed-marker", handle.read()) +class TestSessionLogSurvivesItsDirectory(unittest.TestCase): + """A log moved into a temporary root_log_dir that is then deleted must not + make every later log call in the process raise FileNotFoundError.""" + + def test_a_deleted_log_directory_never_fails_the_caller(self): + root = tempfile.mkdtemp(prefix="wl-log-gone-") + self.addCleanup(shutil.rmtree, root, True) + path = os.path.join(root, "weightslab_logs", "session.log") + os.makedirs(os.path.dirname(path)) + handler = logs._make_file_handler(path) + self.addCleanup(handler.close) + logger = logging.getLogger("test.log_dir_gone") + logger.addHandler(handler) + self.addCleanup(logger.removeHandler, handler) + logger.propagate = False + self.addCleanup(setattr, logger, "propagate", True) + logger.setLevel(logging.INFO) + handler.setLevel(logging.INFO) + + logger.info("before") + handler.close() # set_log_directory / dictConfig close it + shutil.rmtree(os.path.dirname(path)) # the test's tearDown removes the dir + + logger.info("after") # used to raise FileNotFoundError here + handler.flush() + with open(path, encoding="utf-8") as fh: + self.assertIn("after", fh.read()) + + if __name__ == "__main__": unittest.main() diff --git a/weightslab/AGENTS.md b/weightslab/AGENTS.md index 1271bf18..3994b6a0 100644 --- a/weightslab/AGENTS.md +++ b/weightslab/AGENTS.md @@ -260,6 +260,77 @@ def hardness(loss_vals, entropy_vals): ... - No custom classifier needed? Pass `loss_shape_signal=` to `wl.write_dataframe(...)` (§3.7) to compute the built-in tag at dump time. +### 3.6b Live 3-D projection (parametric UMAP) — beta + +On by default, no code: `wl.watch_or_edit(model, flag="model")` hooks the +**input of the last parameterised layer** (last `nn.Linear`, else last conv — +i.e. the representation feeding the head, never its class-logit output), trains a small `features → R³` +encoder on the UMAP objective alongside your model, and writes +`signals//umap_{x,y,z}` per sample. Features are detached and the encoder has +its own optimizer — it cannot perturb training. Studio shows a **Projection** +board (3-D scatter, zoom = server-side level-of-detail, click/lasso selection → +`sample_id` filter on the data board). + +- Off: `WEIGHTSLAB_PROJECTION=0` (removes the hook; zero cost; wins over any + keyword), or `wl.watch_or_edit(model, flag="model", projection=False)` per model + (also removes an earlier model's hook — one projection per process). +- Steer: `projection={"layer": "backbone.fc", "every_n_steps": 20, "min_dist": 0.05}`; + a string is shorthand for the layer (`projection="backbone.fc"`). Keys: + `layer, every_n_steps, out_dim, n_neighbors, graph_size, inner_steps, lr, + min_dist, spread, repulsion_strength, signal_prefix` — each overrides its env var. +- **Bring your own** (t-SNE, umap-learn, PCA…): `wl.save_projection_coords(coords, + batch_ids=ids, prefix="tsne")` with `(N, 2|3)` coords, or let WeightsLab collect + the features: `wl.project_dataset(model, loader, method=TSNE(n_components=3))` + (`method` = anything with `fit_transform`, or a callable `(N,F) -> (N,2|3)`; prefix + defaults to the method's name). Any time — e.g. at `model.get_age() == 500`. + Refused: the prefix `umap` while the built-in runs, shapes other than 2/3 columns. + Not counted as "seen" (no `nb_seen`/`last_seen` bump). The board's picker lists it. + The picker shows only REGISTERED prefixes (`/projection/prefixes.json`, + written by every WeightsLab projection writer), so hand-written `tsne_x` signals via + `save_signals` do not appear — use `save_projection_coords`. No registry file + (an older run) ⇒ every `_x`/`_y` column pair is offered. +- Live coverage: every training batch is buffered (graph = the last `graph_size` + samples) and placed with the current encoder between fits, so all samples get + coordinates each epoch. Nothing is drawn before the first fit. +- From the Studio notebook: a cell can read the live model and call + `wl.save_projection_coords`; running a cell PAUSES training first (and it stays + paused until Play). From the script: `wl.pause_training()` (e.g. at convergence). + Examples: `examples/Usecases/wl-fashion-mnist-umap` (built-in, live), + `examples/Usecases/wl-fashion-mnist-custom-projection` (converge → pause from code + → notebook-written t-SNE → board) and `examples/PyTorch/wl-fashion-mnist-umap` + (1k steps with the live UMAP → reload the checkpoint → discard the lowest-loss + samples → your own 2-D PCA → compare). + `project_dataset` sweeps the WHOLE loader (a tracked loader mid-epoch is restarted), + never ages the model, and returns `sample_ids` aligned with the features your `method` got. +- Studio: with nothing to draw the board shows an orange "No projection found" + ribbon with the backend's reason (not fitted yet / stopped / off / other prefixes). +- Checkpoints: the encoder is saved beside every checkpoint + (`models//projection/.pt`) and restored WITH it — on a restart's + auto-load too, even though the projection attaches after the weights load. A + checkpoint older than the first fit resets the layout; an architecture restore + re-hooks the same layer (by name) on the new model object. +- Failure contract: never raises into training. A failed fit warns once and + retries; 5 in a row turn it off for the run (`stats()["disabled_reason"]`); + NaN/inf features are skipped (encoder never goes NaN); bad options → warning, + model still wrapped. +- Works for every supervised task — it only needs a forward pass plus a + per-sample write carrying `batch_ids` (a `flag="loss"` criterion **or** a + direct `wl.save_signals`). Detection/segmentation `per_instance=True` is + handled: the projection gets the sample-level ids, one point per image. +- Eval batches are placed but never fit on (no eval leakage into the layout). +- Fully dynamic: layer, reduction, feature width and device are all discovered + from the first batch. A live architecture edit (prune/grow) rebuilds the + encoder — the cloud re-lays-out, `stats()["rebuilds"]` counts it. +- Debug: `wl.projection.get_tracker().stats()` (`enabled, layer, fits, failures, + skipped_nonfinite, disabled_reason, …`). `fits == 0` while training means + the auto-picked layer never ran — pass `projection={"layer": ...}`. +- **Wrong layer after the fact?** `wl.project_dataset(model, loader, layer="backbone.layer3", + prefix="umap_layer3", epochs=20)` re-projects offline from a restored checkpoint — + no training step, weights/mode untouched. Different `prefix` values coexist as + separate columns; the board's picker (fed by `available_prefixes`) switches between them. +- Full page: `weightslab/docs/projection.rst`; example: + `examples/Usecases/wl-parametric-umap/`. + ### 3.7 Persisting & inspecting history - `wl.write_history()` / `wl.write_dataframe(path=, format="csv", columns=[...], loss_shape_signal=)` — @@ -321,6 +392,7 @@ Authoritative reference: `weightslab/docs/configuration.rst`. High-signal ones: | `GRPC_AUTH_TOKEN` | unset | Optional token auth on top of mTLS. | | `GRPC_MAX_MESSAGE_BYTES` | `268435456` | Raise if large tensors/images fail to transfer. | | `WEIGHTSLAB_DISABLE_WATCHDOGS` | `0` | Set `1` when breakpoint-debugging (§5). | +| `WEIGHTSLAB_PROJECTION` | on | Live 3-D parametric-UMAP projection (beta). `0`/`false`/`no`/`off` removes the hook entirely. Also `_EVERY` (50), `_DIM` (3), `_NEIGHBORS` (15), `_GRAPH` (512, kNN graph buffer). | | `GRPC_WATCHDOG_STUCK_SECONDS` | `60` | Lock/RPC stuck threshold + lock-acquire timeout. | **Frontend — runtime `window.*` globals (injected at `weightslab start` time; restart+reload to apply):** @@ -399,6 +471,7 @@ around here. - `components/` — `global_monitoring.py` (locks, `guard_*`, pause), `checkpoint_manager.py`, `evaluation_controller.py`. - `data/` — `dataframe_manager.py`, `data_samples_with_ops.py`, `sample_stats.py`, H5 storage (`h5_dataframe_store.py`, `h5_array_store.py`, `array_proxy.py`). - `backend/` — `ledgers.py` (`GLOBAL_LEDGER`), `logger.py`, `audit_logger.py`, `cli.py`. +- `projection/parametric_umap.py` — live 3-D projection (§3.6b): feature hook, UMAP encoder, coordinate write-back, checkpoint pairing; `projection/custom.py` — user projections (`save_projection_coords`); served by `trainer/services/projection_service.py` (`GetProjection`). - `security/` (`CertAuthManager`), `proto/`, `docs/`. - `integrations/ultralytics/` — `WLAwareTrainer`/`WLAwareSegmentationTrainer` (§3.9) — the only wiring done outside a user script. - `examples/{PyTorch,Lightning,Ultralytics,Usecases}//main.py` (+`config.yaml`, `utils/`) — the integration cookbook (§3); `examples/Notebooks/` mirrors most as notebooks, same wiring; `examples/utils/baseline_models/` is a plain model zoo (only the two `wl-classification` examples use it). @@ -407,6 +480,8 @@ around here. - `main.ts` — bootstrap, builds grpc-web transport from `WS_SERVER_*`. - `experiment_service.client.ts`/`experiment_service.ts` — generated client (regen via `npm run generate-proto:data`; don't hand-edit). - `grid_data/` — grid/modal rendering (`GridCell.ts`, `DataImageService.ts`, `gridDataManager.ts`, `BboxRenderer.ts`, `SegmentationRenderer.ts`, `PointCloudViewer.ts`). +- `projection/` — the Projection board (§3.6b): `ProjectionViewer.ts` (three.js scatter, raycast + lasso picking), `projectionBoard.ts` (toolbar, zoom→refetch); `main_area/projectionResizer.ts` sizes it as a third board. +- `main_area/boardRegistry.ts` + `boardLayout.ts` — the main-area boards. **Adding a board = one `BOARDS` entry** (id, label, selector, title selector) plus a header carrying that title element; drag-to-reorder, stacking and side-by-side rows all come for free. Layout is `string[][]` (rows of board ids) in `main-area-layout-v2`; a row with one board is left unwrapped so the per-board resizers keep working, multi-board rows get a `.board-row` flex wrapper. `boardVerticalResizers.ts` still owns per-board heights/collapse; its old two-section DnD is superseded and unwired. - `ui/` — `server.py` (HTTP + gRPC-Web proxy), `static/` (bundled SPA), `utils/` (cert scripts, sync-frontend helper). **Docs:** `weightslab/docs/` — `configuration.rst`, `weights_studio.rst`, `quickstart.rst`, `grpc/`. diff --git a/weightslab/VERSION b/weightslab/VERSION new file mode 100644 index 00000000..b92d1288 --- /dev/null +++ b/weightslab/VERSION @@ -0,0 +1 @@ +2.0.3.dev6 diff --git a/weightslab/__init__.py b/weightslab/__init__.py index fac72cf4..379f3996 100644 --- a/weightslab/__init__.py +++ b/weightslab/__init__.py @@ -55,7 +55,7 @@ } # Everything re-exported straight from .src (attribute name == export name). for _name in ( - "watch_or_edit", "start_training", "serve", "keep_serving", "save_signals", + "watch_or_edit", "start_training", "pause_training", "serve", "keep_serving", "save_signals", "save_instance_signals", "save_group_signals", "save_media", "tag_samples", "save_model_signals", "track_model_signals", "register_categorical_tag", "set_categorical_tag", "discard_samples", @@ -70,7 +70,8 @@ "auto_loss_shape_signal_names", "signal_classifier", "resolve_signal_classifier", "LOSS_SHAPES", "get_current_experiment_hash", "pointcloud_thumbnail", - "pointcloud_boxes", "ai_report_generation", + "pointcloud_boxes", "ai_report_generation", "project_dataset", + "save_projection_coords", ): _LAZY_EXPORTS[_name] = (".src", _name) del _name @@ -173,7 +174,9 @@ def _autoinstall_opencode_on_import(): # why the banner used to show an old version. Preference order: # 1. live git checkout -> derive from the current (or nearest older) git tag # 2. built/installed pkg -> the generated _version.py, else dist metadata -# 3. last resort -> a UTC timestamp (keeps import from ever failing) +# 3. committed release -> weightslab/VERSION, the last release tag, written +# onto main/dev by release.yml (sync-version-file) +# 4. last resort -> a UTC timestamp (keeps import from ever failing) def _resolve_version() -> str: pkg_dir = os.path.dirname(os.path.abspath(__file__)) repo_root = os.path.dirname(pkg_dir) @@ -221,7 +224,18 @@ def _clean(v: str) -> str: except Exception: pass - # 3. Never let version resolution break the import. + # 3. Committed release file: reached only with no git, no build and no + # install (e.g. a source copy without .git), so it never masks the more + # precise versions above — it is the last release, not this commit. + try: + with open(os.path.join(pkg_dir, "VERSION"), encoding="utf-8") as fh: + _v = fh.read().strip() + if _v: + return _v + except Exception: + pass + + # 4. Never let version resolution break the import. from datetime import datetime return datetime.utcnow().strftime("%Y%m%d%H%M%S") @@ -272,6 +286,7 @@ def _clean(v: str) -> str: "guard_testing_context", "ledger", "start_training", + "pause_training", "register_categorical_tag", "set_categorical_tag", "get_current_experiment_hash", @@ -283,6 +298,8 @@ def _clean(v: str) -> str: "write_dataframe", "export_annotations", "ai_report_generation", + "project_dataset", + "save_projection_coords", "classify_loss_shape", "write_loss_shapes", "write_signal_shapes", diff --git a/weightslab/backend/cli.py b/weightslab/backend/cli.py index 11ccbeb7..3ca0ef01 100644 --- a/weightslab/backend/cli.py +++ b/weightslab/backend/cli.py @@ -1147,36 +1147,38 @@ def _handle_command(cmd: str) -> Any: def _client_handler(conn: socket.socket, addr): with conn: - conn_file = conn.makefile('rwb') - # greet - greeting = {'ok': True, 'welcome': 'weightslab-cli', 'time': time.time()} - conn_file.write((json.dumps(greeting) + "\n").encode('utf8')) - conn_file.flush() - while True: - try: + try: + conn_file = conn.makefile('rwb') + # greet + greeting = {'ok': True, 'welcome': 'weightslab-cli', 'time': time.time()} + conn_file.write((json.dumps(greeting) + "\n").encode('utf8')) + conn_file.flush() + while True: line = conn_file.readline() - except ConnectionResetError: - break - if not line: - break - cmd = line.decode('utf8').rstrip('\n') - if cmd.lower().strip() in ('quit', 'exit'): - resp = {'ok': True, 'bye': True} - conn_file.write((json.dumps(resp) + "\n").encode('utf8')) - conn_file.flush() - break - resp = _handle_command(cmd) - # sanitize response to ensure JSON serialization succeeds - try: - safe = _sanitize_for_json(resp) - conn_file.write((json.dumps(safe) + "\n").encode('utf8')) - except Exception: - # fallback: send a stringified representation + if not line: + break + cmd = line.decode('utf8').rstrip('\n') + if cmd.lower().strip() in ('quit', 'exit'): + resp = {'ok': True, 'bye': True} + conn_file.write((json.dumps(resp) + "\n").encode('utf8')) + conn_file.flush() + break + resp = _handle_command(cmd) + # sanitize response to ensure JSON serialization succeeds try: - conn_file.write((json.dumps({'ok': False, 'error': str(resp)}) + "\n").encode('utf8')) + safe = _sanitize_for_json(resp) + conn_file.write((json.dumps(safe) + "\n").encode('utf8')) except Exception: - conn_file.write((json.dumps({'ok': False, 'error': 'unserializable_response'}) + "\n").encode('utf8')) - conn_file.flush() + # fallback: send a stringified representation + try: + conn_file.write((json.dumps({'ok': False, 'error': str(resp)}) + "\n").encode('utf8')) + except Exception: + conn_file.write((json.dumps({'ok': False, 'error': 'unserializable_response'}) + "\n").encode('utf8')) + conn_file.flush() + except ConnectionError: + # The client went away mid-exchange: reset, aborted (WinError + # 10053 on Windows) or a broken pipe. Nothing left to answer. + pass def _server_loop_sock(srv: socket.socket): @@ -1187,7 +1189,15 @@ def _server_loop_sock(srv: socket.socket): _server_host = srv.getsockname()[0] _server_port = srv.getsockname()[1] while True: - conn, addr = srv.accept() + try: + conn, addr = srv.accept() + except OSError as exc: + # Closing the socket is how the server is stopped (main()'s + # Ctrl+C path, tests): accept() then fails with EBADF, or + # EINVAL after a shutdown(). That is the end of serving, not a + # crash -- it used to escape as an unhandled thread exception. + logger.debug("cli_server_socket_closed", extra={"error": str(exc)}) + break t = threading.Thread(target=_client_handler, args=(conn, addr), name='WL-CLI_serving_loop', daemon=True) t.start() finally: diff --git a/weightslab/backend/ledgers.py b/weightslab/backend/ledgers.py index c126b405..7415350c 100644 --- a/weightslab/backend/ledgers.py +++ b/weightslab/backend/ledgers.py @@ -15,6 +15,7 @@ import weakref import logging import os +import sys import time import yaml from collections.abc import MutableMapping @@ -78,6 +79,30 @@ def _plain_get_value(value: Any) -> Any: return _KEEP_AS_PROXY +# `is_a_loop` on a loader says that some `for batch in proxy:` is iterating it, +# so the epoch must END (StopIteration) rather than silently restart. It is a +# count of such loops, not a flag: the training script's own `for batch in +# test_loader:` and an evaluation the Studio triggered on that same loader can +# be in flight together (the first paused mid-epoch). With a flag, the loop that +# finished first cleared it under the other, which then never saw its epoch end +# and ran forever: training froze after a test-split evaluation. +_FOR_LOOPS_LOCK = threading.RLock() + + +def _for_loop_started(it) -> None: + with _FOR_LOOPS_LOCK: + count = getattr(it, "_active_for_loops", 0) + 1 + it._active_for_loops = count + it.is_a_loop = True + + +def _for_loop_ended(it) -> None: + with _FOR_LOOPS_LOCK: + count = max(0, getattr(it, "_active_for_loops", 1) - 1) + it._active_for_loops = count + it.is_a_loop = count > 0 + + class Proxy: """A small forwarding proxy that holds a mutable reference to an object. @@ -555,17 +580,27 @@ def __iter__(self): class _ProxyIterator: def __init__(self, it): self._it = it - self._it.is_a_loop = True + self._ended = False + _for_loop_started(it) def __iter__(self): return self + def close(self): + """Release this loop's hold on the loader, once: at its + StopIteration, or when it is abandoned or stopped early.""" + with _FOR_LOOPS_LOCK: + if self._ended: + return + self._ended = True + _for_loop_ended(self._it) + def __next__(self): try: return next(self._it) except StopIteration: # Let StopIteration propagate naturally - self._it.is_a_loop = False # Loop ends here + self.close() # Loop ends here raise except KeyError: # Quiet by default; only surface this diagnostic when the user @@ -576,6 +611,7 @@ def __next__(self): "KeyError during Proxy iteration. This may indicate the underlying object was modified during iteration. Returning StopIteration to end iteration gracefully." + "\nOtherwise there is a missmatch between data metadata returned, e.g., some metadata has augmentation parameters and other not. Please initialize all metadata with the same keys and types to avoid this error." ) + self.close() raise StopIteration def __del__(self): @@ -588,7 +624,7 @@ def __del__(self): # propagates StopIteration once per epoch instead of # transparently auto-resetting. try: - self._it.is_a_loop = False + self.close() except Exception: pass return _ProxyIterator(underlying_iter) @@ -1286,6 +1322,15 @@ def clear(self) -> None: except Exception: pass + # src caches the dataframe handle, which is one of the proxies just + # dropped: the next registration makes a NEW proxy, and save_signals + # kept writing to the dead one while the new run's table stayed empty. + # wl.clear_all() resets it too; this covers ledgers.clear_all() called + # directly. sys.modules, not an import: never loaded, nothing cached. + src = sys.modules.get("weightslab.src") + if src is not None: + src.DATAFRAME_M = None + def snapshot(self) -> Dict[str, List[str]]: """Return the current keys for all registries (a lightweight snapshot).""" with self._lock: diff --git a/weightslab/components/checkpoint_manager.py b/weightslab/components/checkpoint_manager.py index 5212f397..2c7dbc38 100644 --- a/weightslab/components/checkpoint_manager.py +++ b/weightslab/components/checkpoint_manager.py @@ -1209,6 +1209,99 @@ def _save_changes( # rest of the state (cheap CHECKPOINT; replaces the old JSON snapshot). self.flush_logger_to_disk() + @staticmethod + def projection_sidecar_path(checkpoint_file) -> "Path": + """Where a checkpoint's projection encoder lives. + + In a ``projection/`` SUBDIRECTORY, not beside the checkpoint. The + weight-checkpoint scanner globs ``_step_*.pt`` over the model + directory, and a sibling named ``.projection.pt`` matches that + glob -- it was picked up as a weight checkpoint, its step failed to + parse, and checkpoint selection returned None. A subdirectory cannot + collide with that glob however the sidecar is named. + """ + from pathlib import Path as _Path + path = _Path(checkpoint_file) + return path.parent / "projection" / (path.stem + ".pt") + + def _save_projection_sidecar(self, checkpoint_file) -> None: + """Write the live projection encoder next to *checkpoint_file*. + + Best-effort and silent on failure: a checkpoint that saved the model is + a successful checkpoint whether or not a projection happened to exist. + """ + try: + from weightslab.projection import encoder_state, get_tracker, save_projection + tracker = get_tracker() + state = encoder_state(tracker) + if state is None: + return + import torch as _th + target = self.projection_sidecar_path(checkpoint_file) + target.parent.mkdir(parents=True, exist_ok=True) + _th.save(state, target) + # Also refresh the run-level copy, which is what a plain restart + # (no explicit checkpoint restore) picks up. + save_projection(tracker) + except Exception as exc: + logger.debug(f"Could not save projection sidecar: {exc}") + + def _load_projection_sidecar(self, checkpoint_file) -> bool: + """Bring the projection encoder to the moment *checkpoint_file* was taken. + + Delegates to :func:`weightslab.projection.restore_from_checkpoint`, + which also covers the start-up order (weights restored before the + projection is attached) and resets the layout when the checkpoint + predates the first fit. True when an encoder was, or will be, restored. + """ + try: + from weightslab.projection import restore_from_checkpoint + return restore_from_checkpoint(self.projection_sidecar_path(checkpoint_file)) + except Exception as exc: + logger.debug(f"Could not load projection sidecar: {exc}") + return False + + def _save_custom_projections_sidecar(self, checkpoint_file) -> None: + """Write the user's own projections (t-SNE, PCA, ...) beside the weights. + + They are plain per-sample columns, not part of the model, so without + this a restore would leave them describing a model state that no longer + exists. Best-effort and silent on failure, like the encoder sidecar. + """ + try: + from weightslab.projection import snapshots + dfm = ledgers.get_dataframe() + snapshots.save_sidecar(checkpoint_file, dfm.get_df_view()) + except Exception as exc: + logger.debug(f"Could not save custom projection sidecar: {exc}") + + def _restore_custom_projections(self, checkpoint_file) -> None: + """Make the run's custom projections the ones saved with *checkpoint_file*. + + Those computed later are cleared and leave the picker. A checkpoint with + no sidecar (written before this existed) changes nothing. + """ + if checkpoint_file is None: + return + try: + from weightslab.projection import snapshots + saved = snapshots.load_sidecar(checkpoint_file) + if saved is None: + return + dfm = ledgers.get_dataframe() + snapshots.apply(saved, dfm.get_df_view()) + except Exception as exc: + logger.debug(f"Could not restore custom projections: {exc}") + + @staticmethod + def _reattach_projection(model) -> None: + """Re-hook the live projection on a model object that replaced the old one.""" + try: + from weightslab.projection import reattach_projection + reattach_projection(model) + except Exception as exc: + logger.debug(f"Could not re-attach projection: {exc}") + def save_model_checkpoint( self, model: Optional[th.nn.Module] = None, @@ -1304,6 +1397,15 @@ def save_model_checkpoint( th.save(checkpoint, checkpoint_file) logger.info(f"Saved model checkpoint: {checkpoint_file.name}") + # The projection encoder rides ALONGSIDE the checkpoint, in a + # sibling file -- never inside it, so a weights file stays a + # weights file. Pairing them per checkpoint is what makes a restore + # coherent: rolling the model back to step 1000 while the encoder + # is still the step-5000 one would place old features through a + # newer map and silently draw a layout that never existed. + self._save_projection_sidecar(checkpoint_file) + self._save_custom_projections_sidecar(checkpoint_file) + # Update manifest with latest weight checkpoint for this experiment if update_manifest: self._update_manifest_weight_checkpoint(self.current_exp_hash, checkpoint_file.name, self._step_counter) @@ -2103,6 +2205,14 @@ def load_checkpoint(self, result['loaded_components'].add('weights') step = result['weights'].get('step', 0) + # Bring the projection back to the same moment as these + # weights (see _save_projection_sidecar). Without it the + # restored model's features would be read through an + # encoder trained on a later one. + if self._load_projection_sidecar(checkpoint_file_to_load): + result['loaded_components'].add('projection') + self._last_loaded_checkpoint_file = checkpoint_file_to_load + # Extract RNG state from model checkpoint if available checkpoint_rng_state = result['weights'].get('rng_state') if checkpoint_rng_state: @@ -2285,6 +2395,7 @@ def load_state( logger.info(f"{'='*60}") # Load checkpoint data + self._last_loaded_checkpoint_file = None # set only if weights load below checkpoint_data = self.load_checkpoint( exp_hash=exp_hash, load_model=load_model, @@ -2315,6 +2426,9 @@ def load_state( # Register in ledger ledgers.register_model(model) + # A NEW model object: the projection's hook is still on the + # old one, which nothing calls any more. + self._reattach_projection(model) # Set Model Training Guard guard_training_context.model = model # Train @@ -2391,6 +2505,7 @@ def load_state( try: model = model_data['model'] ledgers.register_model(model) + self._reattach_projection(model) weights = checkpoint_data['weights'] if model and weights and 'model_state_dict' in weights: model.load_state_dict(weights['model_state_dict']) @@ -2540,6 +2655,10 @@ def load_state( # Last, so it reads a dataframe that already has the snapshot applied and # a logger history that is done loading. self._rewind_sample_state(exp_hash, applied_step) + # And the user's own projections, after the rewind: it blanks every + # signal column of the samples it rolls back, custom coordinates + # included, so what the checkpoint saved is put back on top of that. + self._restore_custom_projections(getattr(self, '_last_loaded_checkpoint_file', None)) # Update current experiment hash after everything is loaded success = len(self.error_loading_checkpoint) == 0 diff --git a/weightslab/components/global_monitoring.py b/weightslab/components/global_monitoring.py index 84627c86..c4621cb1 100644 --- a/weightslab/components/global_monitoring.py +++ b/weightslab/components/global_monitoring.py @@ -162,8 +162,7 @@ def resume(self, force: bool = False) -> bool: logger.info('\nAttempting to resume training...') # On resume, first dump any pending changes to checkpoint manager - if self.checkpoint_manager == None: - self.checkpoint_manager = get_checkpoint_manager() + self._get_checkpoint_manager() if self.checkpoint_manager != None: self.checkpoint_manager.update_experiment_hash(first_time=True) self.checkpoint_manager.save_pending_changes() # Write pending change to disk @@ -187,8 +186,12 @@ def is_paused(self): return not self._event.is_set() def _get_checkpoint_manager(self): - if self.checkpoint_manager is None: - self.checkpoint_manager = get_checkpoint_manager() + # The manager the ledger holds NOW, never a cached one: clear_all() (a + # new experiment in the same process, or simply the next test) + # registers a fresh manager behind a fresh proxy, and the old reference + # kept dumping this experiment's checkpoints into the previous one's + # manifest -- where a reload by the current hash could not find them. + self.checkpoint_manager = get_checkpoint_manager() def _is_hash_computed(self): self._get_checkpoint_manager() diff --git a/weightslab/data/data_utils.py b/weightslab/data/data_utils.py index 05ac8446..94e3c81c 100644 --- a/weightslab/data/data_utils.py +++ b/weightslab/data/data_utils.py @@ -591,7 +591,7 @@ def load_raw_image(dataset, index, slice_idx: int = None, rank: int = 0) -> Imag if np_img.ndim == 2: # Grayscale 2D np_img = to_uint8(np_img) - return Image.fromarray(np_img, mode="L") + return Image.fromarray(np_img) elif np_img.ndim == 3: # 3D array - detect channel order if _detect_channel_first_3d(np_img.shape): @@ -602,11 +602,11 @@ def load_raw_image(dataset, index, slice_idx: int = None, rank: int = 0) -> Imag # Now must be (H, W, C) channels = np_img.shape[-1] if channels == 1: - return Image.fromarray(np_img[..., 0], mode="L") + return Image.fromarray(np_img[..., 0]) elif channels == 3: - return Image.fromarray(np_img, mode="RGB") + return Image.fromarray(np_img) elif channels == 4: - return Image.fromarray(np_img, mode="RGBA") + return Image.fromarray(np_img) else: raise ValueError(f"Unsupported channel count: {channels}") else: @@ -633,6 +633,15 @@ def load_raw_image_array(dataset, index, rank: int = 0) -> tuple: """ wrapped = getattr(dataset, "wrapped_dataset", dataset) + # Cheapest path first: a video dataset that can seek a single frame should + # never be asked to decode a whole clip for a thumbnail nobody plays. Only + # datasets defining get_poster_frame() take this branch. + from weightslab.data.video_utils import load_poster_frame_direct + poster = load_poster_frame_direct(wrapped, index) + if poster is not None: + poster_pil = Image.fromarray(poster) + return poster, False, tuple(poster.shape), poster_pil + if hasattr(wrapped, '__getitem__'): np_img, is_volumetric, original_shape = _get_image_array_and_metadata(wrapped, index, rank=rank) @@ -695,20 +704,20 @@ def load_raw_image_array(dataset, index, rank: int = 0) -> tuple: # Convert middle slice to PIL Image for thumbnail if middle_slice.ndim == 2: middle_slice_uint8 = to_uint8(middle_slice) - middle_pil = Image.fromarray(middle_slice_uint8, mode="L") + middle_pil = Image.fromarray(middle_slice_uint8) else: if middle_slice.ndim == 3 and _detect_channel_first_3d(middle_slice.shape): middle_slice = np.transpose(middle_slice, (1, 2, 0)) middle_slice_uint8 = to_uint8(middle_slice) channels = middle_slice_uint8.shape[-1] if middle_slice_uint8.ndim == 3 else 1 if channels == 1: - middle_pil = Image.fromarray(middle_slice_uint8[..., 0], mode="L") + middle_pil = Image.fromarray(middle_slice_uint8[..., 0]) elif channels == 3: - middle_pil = Image.fromarray(middle_slice_uint8, mode="RGB") + middle_pil = Image.fromarray(middle_slice_uint8) elif channels == 4: - middle_pil = Image.fromarray(middle_slice_uint8, mode="RGBA") + middle_pil = Image.fromarray(middle_slice_uint8) else: - middle_pil = Image.fromarray(middle_slice_uint8[..., 0], mode="L") # Fallback + middle_pil = Image.fromarray(middle_slice_uint8[..., 0]) # Fallback return np_img, is_volumetric, original_shape, middle_pil diff --git a/weightslab/data/dataframe_manager.py b/weightslab/data/dataframe_manager.py index f8ba4bec..4e20850f 100644 --- a/weightslab/data/dataframe_manager.py +++ b/weightslab/data/dataframe_manager.py @@ -880,8 +880,6 @@ def _load_existing_data(self, origin: str = None, autoload_arrays: bool | list | # Keep the (sample_id, annotation_id) index clean and grouped. if self._df.index.has_duplicates: self._df = self._df[~self._df.index.duplicated(keep='last')] - if isinstance(self._df.index, pd.MultiIndex): - self._df = self._df.sort_index() else: logger.warning(f"[LedgeredDataFrameManager] Loaded data missing 'sample_id' column for origin={origin}. Skipping load.") self._apply_recovery_snapshot(origin) diff --git a/weightslab/data/h5_array_store.py b/weightslab/data/h5_array_store.py index abb155c5..0c137e28 100644 --- a/weightslab/data/h5_array_store.py +++ b/weightslab/data/h5_array_store.py @@ -167,7 +167,7 @@ def release_read(self): try: self._readers -= 1 if self._readers == 0: - self._read_ready.notifyAll() + self._read_ready.notify_all() finally: self._read_ready.release() @@ -190,7 +190,7 @@ def release_write(self): self._read_ready.acquire() try: self._writers -= 1 - self._read_ready.notifyAll() + self._read_ready.notify_all() finally: self._read_ready.release() diff --git a/weightslab/data/point_cloud_utils.py b/weightslab/data/point_cloud_utils.py index dc090ae8..ccbb8a07 100644 --- a/weightslab/data/point_cloud_utils.py +++ b/weightslab/data/point_cloud_utils.py @@ -351,7 +351,7 @@ def point_cloud_to_bev_image( pts = filter_valid_points(points) if pts.shape[0] == 0: - return Image.fromarray(canvas, mode="RGB") + return Image.fromarray(canvas) if pc_range is None: pc_range = get_pc_range(None, pts) @@ -363,7 +363,7 @@ def point_cloud_to_bev_image( ) pts = pts[in_range] if pts.shape[0] == 0: - return Image.fromarray(canvas, mode="RGB") + return Image.fromarray(canvas) # Pixel coordinates: col follows +x (right), row follows -y (so +y is up). cols = ((pts[:, 0] - x_min) / max(x_max - x_min, 1e-6) * (size - 1)).astype(np.int64) @@ -389,7 +389,7 @@ def point_cloud_to_bev_image( flat_canvas = canvas.reshape(-1, 3) flat_canvas[occupied] = colors - return Image.fromarray(canvas, mode="RGB") + return Image.fromarray(canvas) def point_cloud_to_range_image( @@ -422,7 +422,7 @@ def point_cloud_to_range_image( pts = filter_valid_points(points) if pts.shape[0] == 0: canvas = np.full((image_height, image_width, 3), np.array(_BEV_BACKGROUND, dtype=np.uint8), dtype=np.uint8) - return Image.fromarray(canvas, mode="RGB") + return Image.fromarray(canvas) # Extract coordinates x, y = pts[:, 0], pts[:, 1] @@ -471,7 +471,7 @@ def point_cloud_to_range_image( colors = np.clip(colors * brightness[:, None], 0, 255).astype(np.uint8) canvas[v, u] = colors - return Image.fromarray(canvas, mode="RGB") + return Image.fromarray(canvas) def boxes_dimensionality(boxes: np.ndarray) -> int: @@ -574,9 +574,7 @@ def render_thumbnail_2d_for_dataset(dataset, points: np.ndarray) -> Image.Image: if isinstance(img, Image.Image): return img.convert("RGB") arr = np.asarray(img) - if arr.ndim == 2: - return Image.fromarray(arr.astype(np.uint8), mode="L").convert("RGB") - return Image.fromarray(arr.astype(np.uint8), mode="RGB") + return Image.fromarray(arr.astype(np.uint8)).convert("RGB") except Exception as exc: logger.warning("Dataset %s hook failed (%s); falling back to default.", THUMBNAIL_2D_HOOK, exc) @@ -588,9 +586,7 @@ def render_thumbnail_2d_for_dataset(dataset, points: np.ndarray) -> Image.Image: if isinstance(img, Image.Image): return img.convert("RGB") arr = np.asarray(img) - if arr.ndim == 2: - return Image.fromarray(arr.astype(np.uint8), mode="L").convert("RGB") - return Image.fromarray(arr.astype(np.uint8), mode="RGB") + return Image.fromarray(arr.astype(np.uint8)).convert("RGB") except Exception as exc: logger.warning("Dataset %s hook failed (%s); falling back to default.", BEV_IMAGE_HOOK, exc) @@ -601,9 +597,7 @@ def render_thumbnail_2d_for_dataset(dataset, points: np.ndarray) -> Image.Image: if isinstance(img, Image.Image): return img.convert("RGB") arr = np.asarray(img) - if arr.ndim == 2: - return Image.fromarray(arr.astype(np.uint8), mode="L").convert("RGB") - return Image.fromarray(arr.astype(np.uint8), mode="RGB") + return Image.fromarray(arr.astype(np.uint8)).convert("RGB") except Exception as exc: logger.warning("Registered thumbnail fn failed (%s); falling back to default.", exc) diff --git a/weightslab/data/video_utils.py b/weightslab/data/video_utils.py index b17d87cb..66ca4beb 100644 --- a/weightslab/data/video_utils.py +++ b/weightslab/data/video_utils.py @@ -67,6 +67,9 @@ FPS_ATTRS = ("fps", "frame_rate") AUDIO_HOOK = "get_audio" POSTER_HOOK = "render_video_poster" +# Optional: return the grid's still directly, so the clip is never decoded +# just to throw all but one frame away (see load_poster_frame_direct). +POSTER_FRAME_HOOK = "get_poster_frame" # Global decorator registry (set via wl.video_poster). _REGISTERED_POSTER_FN = None @@ -104,6 +107,112 @@ def is_audio_task(task_type) -> bool: return str(task_type or "").strip().lower() == AUDIO_GENERATION_TASK +def load_poster_frame_direct(dataset, index): + """Ask the dataset for a poster frame WITHOUT decoding the whole clip. + + The grid draws one frame per sample, but the only way to get it used to be + to load the sample — which for a video dataset means decoding every frame + of the clip just to throw all but one away. Measured on a 40 s 720p corpus: + 283 ms to decode a 16-frame window against 158 ms for a single frame, so + roughly half the grid's image time was spent on frames nobody looks at. + + A dataset opts in by defining ``get_poster_frame(index)`` returning an + ``[H, W, C]`` array or a PIL image (None to decline, e.g. when it cannot + cheaply seek). Datasets that don't are unaffected — the caller falls back + to decoding the clip. + """ + wrapped = getattr(dataset, "wrapped_dataset", dataset) + for source in (wrapped, dataset): + hook = getattr(source, POSTER_FRAME_HOOK, None) + if not callable(hook): + continue + try: + frame = hook(int(index)) + except Exception as exc: + logger.warning("Dataset %s hook failed: %r", POSTER_FRAME_HOOK, exc) + return None + if frame is None: + return None + if isinstance(frame, Image.Image): + return np.asarray(frame) + arr = np.asarray(frame) + # Reject anything that is not a single frame rather than guessing: + # a wrong-shaped poster would be drawn, not raise. + if arr.ndim == 2 or (arr.ndim == 3 and arr.shape[-1] in _VIDEO_CHANNELS): + return _frames_to_uint8(arr[None])[0] + logger.warning( + "%s returned shape %s, expected [H, W] or [H, W, C]; ignoring", + POSTER_FRAME_HOOK, arr.shape) + return None + return None + + +def poster_frame_index(num_frames: int) -> int: + """Index of the frame :func:`video_poster_frame` renders. + + Kept as its own function because anything drawn ON the poster — a + segmentation mask, a set of boxes — has to be taken from the same frame, + or the overlay describes a moment the picture does not show. + """ + return max(0, int(num_frames) // 2) + + +def select_frame_annotation(annotation, num_frames: int): + """Reduce a per-frame annotation to the one the poster shows. + + Video segmentation hands us a ``[T, H, W]`` mask and video detection a + length-T sequence of box arrays, but the grid and the modal's still view + draw exactly one frame. Encoding the whole stack instead would RLE a 3-D + array into a stat the renderer reads as 2-D, i.e. garbage. + + Anything that is not clearly per-frame is returned untouched, so image + datasets and single-mask video datasets are unaffected. + + Args: + annotation: ``[T, H, W]`` / ``[T, H, W, C]`` array, or a sequence of + length T whose elements are that frame's annotation. + num_frames: T, the clip's frame count. + + Returns: + The annotation for the poster frame, or ``annotation`` unchanged. + """ + if annotation is None or num_frames is None or int(num_frames) <= 1: + return annotation + + index = poster_frame_index(num_frames) + + shape = getattr(annotation, "shape", None) + if shape is not None: + # A [T, H, W] mask; [H, W] stays as it is. A 4-D [T, H, W, C] is also + # per-frame, but a 3-D [H, W, C] colour mask must not be sliced, hence + # the leading-dimension check against T rather than rank alone. + if len(shape) >= 3 and shape[0] == int(num_frames): + return annotation[index] + return annotation + + if isinstance(annotation, (list, tuple)) and len(annotation) == int(num_frames): + return annotation[index] + return annotation + + +def has_playable_media(dataset, task_type=None) -> bool: + """True when GetMedia can stream this sample's own content. + + Deliberately wider than :func:`is_video_task`: a clip is playable because + of what it *is*, not because of what the model does with it. A video + classification or detection dataset declares ``media_kind = "video"`` and + keeps its own task type, and must still get the grid badge and the modal + player — which is exactly what this module's header promises. + + Pair it with :func:`is_video_sample` (which additionally inspects the + array) wherever the frames are actually in hand; this predicate is for + the metadata paths, where they are not. + """ + if is_video_task(task_type) or is_audio_task(task_type): + return True + return get_media_kind(dataset) in ("video", "audio") + + def get_media_kind(dataset) -> str: """Return the dataset's declared media kind ("" when it declares none).""" wrapped = getattr(dataset, "wrapped_dataset", dataset) @@ -211,15 +320,15 @@ def video_frame_at(frames, index: int) -> Image.Image: def _frame_to_pil(frame: np.ndarray) -> Image.Image: """One [H, W, C] uint8 frame -> PIL image.""" if frame.ndim == 2: - return Image.fromarray(frame, mode="L") + return Image.fromarray(frame) channels = frame.shape[-1] if channels == 1: - return Image.fromarray(frame[..., 0], mode="L") + return Image.fromarray(frame[..., 0]) if channels == 4: - return Image.fromarray(frame, mode="RGBA") + return Image.fromarray(frame) if channels == 3: - return Image.fromarray(frame, mode="RGB") - return Image.fromarray(frame[..., 0], mode="L") + return Image.fromarray(frame) + return Image.fromarray(frame[..., 0]) def video_poster_frame(dataset, frames) -> Image.Image: @@ -399,16 +508,35 @@ def encode_video_mp4(frames, fps: float, audio=None, sample_rate: int = 0, "-movflags", "+faststart", out_path, ] - proc = subprocess.run( - cmd, input=arr.tobytes(), stdout=subprocess.PIPE, - stderr=subprocess.PIPE, check=False) - if proc.returncode != 0 or not os.path.exists(out_path): + payload = arr.tobytes() + # One retry. The first exec of the ffmpeg binary in a process can fail + # on Windows (AV/SmartScreen inspecting it, or a slow package shim) + # with a non-zero status and NO stderr at all — observed once, on the + # first GetMedia of a run. Without a retry that transient downgrades + # the clip to GIF, and because the encoded clip is then cached, the + # sample stays GIF for the life of the process. + for attempt in (1, 2): + proc = subprocess.run( + cmd, input=payload, stdout=subprocess.PIPE, + stderr=subprocess.PIPE, check=False) + if proc.returncode == 0 and os.path.exists(out_path): + with open(out_path, "rb") as handle: + return handle.read() + + stderr = proc.stderr.decode("utf-8", "replace").strip() logger.warning( - "ffmpeg failed to encode %d frames: %s", - frame_count, proc.stderr.decode("utf-8", "replace")[:400]) - return b"" - with open(out_path, "rb") as handle: - return handle.read() + "ffmpeg attempt %d/2 failed to encode %d frames " + "(returncode=%s, output_exists=%s, %d input bytes): %s", + attempt, frame_count, proc.returncode, + os.path.exists(out_path), len(payload), + stderr[:400] or "") + if attempt == 1: + # Do not leave a partial file behind for the exists() check. + try: + os.remove(out_path) + except OSError: + pass + return b"" except Exception as exc: logger.warning("MP4 encoding failed: %r", exc) return b"" diff --git a/weightslab/docker/docker/utils/test-deployment.sh b/weightslab/docker/docker/utils/test-deployment.sh deleted file mode 100644 index c58487b6..00000000 --- a/weightslab/docker/docker/utils/test-deployment.sh +++ /dev/null @@ -1,437 +0,0 @@ -#!/bin/bash -# Comprehensive deployment test for weights_studio frontend + backend -# Tests communication between frontend and Python backend (wl-segmentation example) -# Runs both unsecured (HTTP) and secured (HTTPS) test scenarios - -# set -e - -# Colors for output -RED='\033[0;31m' -GREEN='\033[0;32m' -YELLOW='\033[1;33m' -BLUE='\033[0;34m' -NC='\033[0m' # No Color - -# Configuration -CURRENT_DIR="$(cd "$(dirname "$0")/../.." && pwd)" -DOCKER_DIR="$CURRENT_DIR/docker" -UTILS_DIR="$DOCKER_DIR/utils" -EXAMPLE_DIR="$CURRENT_DIR/tests/playwright/examples/wl-segmentation" -GRPC_PORT=50051 -ENVOY_PORT=8080 -FRONTEND_PORT=5173 -HEALTH_CHECK_TIMEOUT=30 -COMMUNICATION_TEST_TIMEOUT=20 - -# Global test results -TESTS_PASSED=0 -TESTS_FAILED=0 - -# ============================================================================ -# Utility Functions -# ============================================================================ - -log_info() { - echo -e "${BLUE}[INFO]${NC} $1" -} - -log_success() { - echo -e "${GREEN}[✓ PASS]${NC} $1" - ((TESTS_PASSED++)) -} - -log_error() { - echo -e "${RED}[✗ FAIL]${NC} $1" - ((TESTS_FAILED++)) -} - -log_warn() { - echo -e "${YELLOW}[WARN]${NC} $1" -} - -print_section() { - echo "" - echo -e "${BLUE}╔════════════════════════════════════════════════════════════════╗${NC}" - echo -e "${BLUE}║${NC} $1" - echo -e "${BLUE}╚════════════════════════════════════════════════════════════════╝${NC}" - echo "" -} - -# Clean up docker resources -cleanup() { - log_info "Cleaning up Docker resources..." - cd "$DOCKER_DIR" - docker compose -f docker-compose.yml -f docker-compose.prod.yml down --remove-orphans 2>/dev/null || true - # Force remove any lingering containers - docker ps -a | grep weights_studio | awk '{print $1}' | xargs -r docker rm -f 2>/dev/null || true - # Remove network - docker network rm docker_weights_studio_network 2>/dev/null || true - sleep 3 -} - -# Wait for service to be ready -wait_for_service() { - local service_name=$1 - local port=$2 - local timeout=$3 - local protocol=${4:-http} - - log_info "Waiting for $service_name on $protocol://localhost:$port..." - - local start_time=$(date +%s) - while true; do - local current_time=$(date +%s) - local elapsed=$((current_time - start_time)) - - if [ $elapsed -gt $timeout ]; then - log_error "$service_name failed to start after ${timeout}s" - return 1 - fi - - if [ "$protocol" = "https" ]; then - if curl -sk -o /dev/null -w "%{http_code}" "https://localhost:$port/health" 2>/dev/null | grep -q "200"; then - log_success "$service_name is ready" - return 0 - fi - else - if curl -s -o /dev/null -w "%{http_code}" "http://localhost:$port/health" 2>/dev/null | grep -q "200"; then - log_success "$service_name is ready" - return 0 - fi - fi - - sleep 2 - done -} - -# Check if gRPC backend is responding -check_grpc_backend() { - local timeout=$1 - log_info "Checking gRPC backend on localhost:$GRPC_PORT..." - - local start_time=$(date +%s) - while true; do - local current_time=$(date +%s) - local elapsed=$((current_time - start_time)) - - if [ $elapsed -gt $timeout ]; then - log_error "gRPC backend not responding after ${timeout}s" - return 1 - fi - - # Try to connect using nc (netcat) - if nc -z localhost $GRPC_PORT 2>/dev/null; then - log_success "gRPC backend is listening on port $GRPC_PORT" - return 0 - fi - - sleep 1 - done -} - -# Check frontend-backend communication -test_frontend_backend_communication() { - local protocol=$1 - local test_name=$2 - local port=${3:-$FRONTEND_PORT} - - log_info "Testing frontend -> backend communication ($test_name)..." - - # The frontend makes requests to /api which proxies to envoy - # Envoy proxies to localhost:50051 (the gRPC backend) - - # Check if frontend can resolve the backend - local health_url="$protocol://localhost:$port/health" - - local http_code=$(curl -sk -o /dev/null -w "%{http_code}" "$health_url" 2>/dev/null || echo "000") - - if [ "$http_code" = "200" ]; then - log_success "Frontend is responding ($test_name)" - else - log_error "Frontend health check failed with HTTP $http_code ($test_name)" - return 1 - fi - - return 0 -} - -# Run backend example in background and monitor -run_backend_example() { - local test_scenario=$1 - - if [ ! -f "$EXAMPLE_DIR/main.py" ]; then - log_error "Python example not found at $EXAMPLE_DIR/main.py" - return 1 - fi - - log_info "Starting backend example: $test_scenario..." - - cd "$EXAMPLE_DIR" - - # Start Python backend in background - timeout 60s python main.py > "/tmp/ws_backend_${test_scenario}.log" 2>&1 & - local backend_pid=$! 2>/dev/null || true - - log_info "Backend process started (PID: $backend_pid)" - - # Wait a bit for backend to initialize - sleep 3 - - # Check if process is still running - if ! kill -0 $backend_pid 2>/dev/null; then - log_error "Backend process exited prematurely" - cat "/tmp/ws_backend_${test_scenario}.log" - return 1 - fi - - log_success "Backend example is running" - - # Check gRPC port - if check_grpc_backend $COMMUNICATION_TEST_TIMEOUT; then - log_success "Backend gRPC service is reachable" - else - log_error "Backend gRPC service is not reachable" - kill $backend_pid 2>/dev/null || true - return 1 - fi - - return 0 -} - -# Temporarily hide/show certs and tokens for testing -hide_certs() { - local certs_dir=$1 - if [ -d "$certs_dir" ]; then - mv "$certs_dir" "${certs_dir}.bak" - mkdir -p "$certs_dir" - fi -} - -restore_certs() { - local certs_dir=$1 - if [ -d "${certs_dir}.bak" ]; then - rm -rf "$certs_dir" - mv "${certs_dir}.bak" "$certs_dir" - fi -} - -hide_grpc_token() { - local token_file=$1 - if [ -f "$token_file" ]; then - mv "$token_file" "${token_file}.bak" - fi -} - -restore_grpc_token() { - local token_file=$1 - if [ -f "${token_file}.bak" ]; then - mv "${token_file}.bak" "$token_file" - fi -} - -# Run test scenario -run_test_scenario() { - local scenario_name=$1 - local secure=$2 - local use_certs=${3:-true} - local use_auth=${4:-false} - - # Create a safe filename from scenario name (remove spaces, colons, etc) - local scenario_safe=$(echo "$scenario_name" | tr '[:upper:]' '[:lower:]' | sed 's/[^a-z0-9]/-/g' | sed 's/-\+/-/g' | sed 's/-$//') - - print_section "$scenario_name" - - # Cleanup previous deployment - cleanup - - # Setup certs and tokens for this scenario - local certs_dir="$DOCKER_DIR/.weightslab-certs-test" - log_info "certs_dir: $certs_dir" - local token_file="$certs_dir/.grpc_auth_token" - local certs_hidden=false - local token_hidden=false - - if [ "$use_certs" = "false" ] && [ -d "$certs_dir" ]; then - log_info "Temporarily hiding TLS certificates for this test..." - hide_certs "$certs_dir" - certs_hidden=true - fi - - if [ "$use_auth" = "false" ] && [ -f "$token_file" ]; then - log_info "Temporarily hiding gRPC auth token for this test..." - hide_grpc_token "$token_file" - token_hidden=true - fi - - # Set environment for deployment - if [ "$secure" = "true" ]; then - if [ "$use_certs" = "true" ] && [ "$use_auth" = "true" ]; then - log_info "Secure mode: HTTPS with TLS + gRPC auth enabled" - elif [ "$use_certs" = "true" ]; then - log_info "Secure mode: HTTPS with TLS (no gRPC auth)" - elif [ "$use_auth" = "true" ]; then - log_info "Secure mode: HTTP with gRPC auth only (no TLS)" - fi - export E2E_MANAGED_STACK=1 - else - log_info "Unsecure mode: HTTP without TLS or gRPC auth" - export E2E_MANAGED_STACK=1 - fi - - # Build and deploy frontend - cd "$UTILS_DIR" - - if [ "$secure" = "true" ] && [ "$use_certs" = "true" ]; then - log_info "Building and deploying with HTTPS support..." - WEIGHTSLAB_CERTS_DIR="$certs_dir" bash build-and-deploy.sh > "/tmp/build_deploy_${scenario_safe}.log" 2>&1 || { - log_error "Build and deploy failed" - cat "/tmp/build_deploy_${scenario_safe}.log" - [ "$certs_hidden" = true ] && restore_certs "$certs_dir" - [ "$token_hidden" = true ] && restore_grpc_token "$token_file" - return 1 - } - else - log_info "Building and deploying with HTTP (no TLS certs)..." - WEIGHTSLAB_CERTS_DIR="$certs_dir" bash build-and-deploy.sh --unsecure > "/tmp/build_deploy_${scenario_safe}.log" 2>&1 || { - log_error "Build and deploy failed" - cat "/tmp/build_deploy_${scenario_safe}.log" - [ "$certs_hidden" = true ] && restore_certs "$certs_dir" - [ "$token_hidden" = true ] && restore_grpc_token "$token_file" - return 1 - } - fi - - log_success "Frontend deployment successful" - - # Wait for services to be ready - local protocol="http" - local frontend_port=$FRONTEND_PORT - if [ "$secure" = "true" ] && [ "$use_certs" = "true" ]; then - protocol="https" - fi - - # Check frontend (don't check envoy since it depends on upstream backend) - if ! wait_for_service "Frontend (nginx)" "$frontend_port" "$HEALTH_CHECK_TIMEOUT" "$protocol"; then - log_error "Frontend failed to start" - [ "$certs_hidden" = true ] && restore_certs "$certs_dir" - [ "$token_hidden" = true ] && restore_grpc_token "$token_file" - return 1 - fi - - # Test frontend-backend communication setup - if test_frontend_backend_communication "$protocol" "$scenario_safe" "$frontend_port"; then - log_success "Frontend-backend communication setup OK" - else - log_error "Frontend-backend communication setup failed" - [ "$certs_hidden" = true ] && restore_certs "$certs_dir" - [ "$token_hidden" = true ] && restore_grpc_token "$token_file" - return 1 - fi - - # Restore certs and tokens - if [ "$certs_hidden" = true ]; then - log_info "Restoring TLS certificates..." - restore_certs "$certs_dir" - fi - if [ "$token_hidden" = true ]; then - log_info "Restoring gRPC auth token..." - restore_grpc_token "$token_file" - fi - - return 0 -} - -# ============================================================================ -# Main Test Execution -# ============================================================================ - -main() { - print_section "WEIGHTS STUDIO FRONTEND-BACKEND DEPLOYMENT TEST" - - log_info "Test Configuration:" - log_info " Weights Studio Dir: $CURRENT_DIR" - log_info " Docker Dir: $DOCKER_DIR" - log_info " Example Dir: $EXAMPLE_DIR" - log_info " Envoy Port: $ENVOY_PORT" - log_info " Frontend Port: $FRONTEND_PORT" - log_info " gRPC Backend Port: $GRPC_PORT" - - # Test A: Unsecured Environment (HTTP) - if run_test_scenario "TEST A: UNSECURED ENVIRONMENT (HTTP)" "false" "false" "false"; then - log_success "TEST A PASSED: Frontend and backend communicate over HTTP" - else - log_error "TEST A FAILED: Frontend-backend communication issue in HTTP mode" - fi - - # Cleanup between tests - cleanup - sleep 5 - - # Test B: Secured Environment (HTTPS with TLS + gRPC Auth) - # Note: This requires TLS certificates AND gRPC auth token to be present - if [ -d "$DOCKER_DIR/.weightslab-certs-test" ] && [ -f "$DOCKER_DIR/.weightslab-certs-test/envoy-server.crt" ]; then - if [ -f "$DOCKER_DIR/.weightslab-certs-test/.grpc_auth_token" ]; then - if run_test_scenario "TEST B: SECURED ENVIRONMENT (HTTPS + TLS + gRPC AUTH)" "true" "true" "true"; then - log_success "TEST B PASSED: Frontend and backend communicate over HTTPS with full security" - else - log_error "TEST B FAILED: Frontend-backend communication issue in HTTPS+Auth mode" - fi - else - log_warn "TEST B SKIPPED: gRPC auth token not found" - log_warn " Token should be at $DOCKER_DIR/.weightslab-certs-test/.grpc_auth_token" - fi - - # Cleanup between tests - cleanup - sleep 5 - - # Test C: TLS Certificates Only (no gRPC auth) - if run_test_scenario "TEST C: SECURED ENVIRONMENT (HTTPS + TLS ONLY, NO GRPC AUTH)" "true" "true" "false"; then - log_success "TEST C PASSED: Frontend and backend communicate over HTTPS (TLS only, no gRPC auth)" - else - log_error "TEST C FAILED: Frontend-backend communication issue in TLS-only mode" - fi - - # Cleanup between tests - cleanup - sleep 5 - - # Test D: gRPC Auth Only (no TLS certificates) - if [ -f "$DOCKER_DIR/.weightslab-certs-test/.grpc_auth_token" ]; then - if run_test_scenario "TEST D: SECURED ENVIRONMENT (HTTP + gRPC AUTH ONLY, NO TLS)" "true" "false" "true"; then - log_success "TEST D PASSED: Frontend and backend communicate over HTTP with gRPC auth only" - else - log_error "TEST D FAILED: Frontend-backend communication issue in gRPC-auth-only mode" - fi - else - log_warn "TEST D SKIPPED: gRPC auth token not found" - log_warn " Token should be at $DOCKER_DIR/.weightslab-certs-test/.grpc_auth_token" - fi - else - log_warn "TESTS B, C, D SKIPPED: TLS certificates not found at $DOCKER_DIR/.weightslab-certs-test" - log_warn " To run secured tests, copy certs to:" - log_warn " $DOCKER_DIR/.weightslab-certs-test" - fi - - # Final cleanup - cleanup - - # Print summary - print_section "TEST SUMMARY" - echo -e "Tests Passed: ${GREEN}$TESTS_PASSED${NC}" - echo -e "Tests Failed: ${RED}$TESTS_FAILED${NC}" - - if [ $TESTS_FAILED -eq 0 ]; then - log_success "ALL TESTS PASSED! ✨" - return 0 - else - log_error "SOME TESTS FAILED ❌" - return 1 - fi -} - -# Handle script interruption -trap cleanup EXIT - -# Run main -main "$@" diff --git a/weightslab/examples/PyTorch/wl-fashion-mnist-umap/config.yaml b/weightslab/examples/PyTorch/wl-fashion-mnist-umap/config.yaml new file mode 100644 index 00000000..b9e08b13 --- /dev/null +++ b/weightslab/examples/PyTorch/wl-fashion-mnist-umap/config.yaml @@ -0,0 +1,49 @@ +# Fashion-MNIST: the built-in live UMAP, then a PCA of your own to compare it with. +experiment_name: fashion_mnist_umap_vs_pca +device: auto +training_steps_to_do: null # the stages below set the schedule, not a ceiling +# root_log_dir: # defaults to $WEIGHTSLAB_ROOT_LOG_DIR, else a temp dir + +optimizer: + lr: 0.001 + +# The built-in projection is ON by default; these keys only tune it. Each one +# overrides its WEIGHTSLAB_PROJECTION_* environment variable. See docs/projection.rst. +projection: + every_n_steps: 25 # fit the UMAP encoder every 25 training steps (default 50) + graph_size: 1024 # samples in the kNN graph (default 512) + n_neighbors: 15 + +# The staged exploration (main.py); `--steps INITIAL RETRAIN` overrides the two step counts. +exploration: + steps_initial: 1000 # train with the live UMAP, then keep a checkpoint + discard_fraction: 0.3 # then drop this share of the training samples, lowest loss first + steps_retrain: 1000 # train on what is left (0 = stop after the comparison) + knn_k: 10 # neighbours counted by the class-purity score + +# Experiment cadence +eval_full_to_train_steps_ratio: 250 +experiment_dump_to_train_steps_ratio: 1000 +skip_checkpoint_load: false +tqdm_display: true +start_timeout: 3 # seconds to let the studio attach before the first step + +# Global dataframe storage +ledger_enable_flushing_threads: true +ledger_enable_h5_persistence: true +ledger_flush_max_rows: 15000 +ledger_flush_interval: 30.0 + +serving_grpc: true + +# data_root: ./data # auto-downloaded on first run +data: + train_loader: + shuffle: true + batch_size: 64 + max_samples: 20000 # raise (or drop) for the full 60k + test_loader: + shuffle: false + batch_size: 256 + max_samples: 2000 + drop_last: false diff --git a/weightslab/examples/PyTorch/wl-fashion-mnist-umap/main.py b/weightslab/examples/PyTorch/wl-fashion-mnist-umap/main.py new file mode 100644 index 00000000..4f69c258 --- /dev/null +++ b/weightslab/examples/PyTorch/wl-fashion-mnist-umap/main.py @@ -0,0 +1,535 @@ +"""Fashion-MNIST classification, explored with projections: the built-in +parametric UMAP first, then a PCA of your own to compare it with. + +A plain PyTorch loop, driven in stages so you can look at the Projection board +between them (Weights Studio -> Data Exploration -> Projection): + +1. **Train** ``steps_initial`` steps (default 1000) with the live 3-D parametric + UMAP running beside the model -- there is no projection code in the loop, + wrapping the model attaches it -- and keep a checkpoint of those weights. +2. **Reload** that checkpoint: weights, optimizer, the UMAP encoder and every + sample's loss as they were at step 1000. (If you let the run go on, this is + the way back to the weights the next stages are about.) +3. **Discard** the ``discard_fraction`` (default 30 %) of training samples with + the lowest loss: the ones the model already gets right, which teach it little. +4. **PCA, 2-D**, written in plain torch below, of what is left -- stored as the + projection ``pca2d_step``. WeightsLab collects the features; + ``wl.project_dataset(model, loader, method=pca_coordinates)`` hands them to your + code. +5. **Compare** the UMAP and the PCA: in the board's picker, and here, with a + side-by-side picture and a number (how often a point's nearest neighbours in + the picture share its class). +6. **Retrain** ``steps_retrain`` steps on what is left; the live UMAP follows + the model, the PCA stays the snapshot you took. + +Run:: + + python main.py # then open the studio (weightslab start) + python main.py --steps 200 100 # a quick pass: 200 initial + 100 retrain + python main.py --wait # stop for Enter before retraining, to look +""" + +import argparse +import logging +import os +import ssl +import tempfile + +# Windows SSL fix: some Windows cert stores contain malformed ASN1 certs that +# crash ssl.create_default_context(). Fall back to unverified only when broken. +try: + ssl.create_default_context() +except ssl.SSLError: + ssl._create_default_https_context = ssl._create_unverified_context + +import numpy as np +import pandas as pd +import torch +import torch.nn as nn +import torch.optim as optim +import tqdm +import yaml +from torch.utils.data import Dataset +from torchmetrics.classification import Accuracy +from torchvision import datasets, transforms + +import weightslab as wl +from weightslab.backend import ledgers + +logging.basicConfig(level=logging.ERROR) +logger = logging.getLogger(__name__) + +CLASS_NAMES = ( + "T-shirt/top", "Trouser", "Pullover", "Dress", "Coat", + "Sandal", "Shirt", "Sneaker", "Bag", "Ankle boot", +) +LOSS_COLUMN = "signals//train-loss-CE" +TRAIN_SPLIT = "train_loader" + + +# ============================================================================= +# Dataset +# ============================================================================= +class FashionMNISTDataset(Dataset): + """Fashion-MNIST yielding ``(image, sample_id, label)``. + + ``sample_id`` is offset per split (``id_base``) so train and test ids never + collide in the shared ledger. + """ + + def __init__(self, root, train=True, download=True, transform=None, + max_samples=None, id_base=0): + self.data = datasets.FashionMNIST(root=root, train=train, + download=download, transform=None) + self.transform = transform + self.train = train + self.max_samples = max_samples + self.id_base = id_base + + def __len__(self): + if self.max_samples is not None: + return min(len(self.data), self.max_samples) + return len(self.data) + + def __getitem__(self, idx): + image, label = self.data[idx] + if self.transform: + image = self.transform(image) + return image, self.id_base + idx, label + + def fast_get_label(self, idx): + """Lets the ledger read labels at init without decoding every image.""" + return int(self.data.targets[idx]) + + def get_metadata(self, idx): + label = int(self.data.targets[idx]) + return {"class_name": CLASS_NAMES[label], + "split": "train" if self.train else "test"} + + +# ============================================================================= +# Model +# ============================================================================= +class FashionCNN(nn.Module): + """Two conv blocks, then a 128-D representation and the classifier. + + ``fc2`` is the last Linear, so the projection hooks ITS input -- the 128-D + output of ``fc1`` after the ReLU, the representation the classifier decides + from (never the logits). + """ + + def __init__(self, num_classes=10): + super().__init__() + self.input_shape = (1, 1, 28, 28) + self.conv1 = nn.Conv2d(1, 32, 3, padding=1) + self.bn1 = nn.BatchNorm2d(32) + self.relu1 = nn.ReLU() + self.pool1 = nn.MaxPool2d(2) # 28 -> 14 + self.conv2 = nn.Conv2d(32, 64, 3, padding=1) + self.bn2 = nn.BatchNorm2d(64) + self.relu2 = nn.ReLU() + self.pool2 = nn.MaxPool2d(2) # 14 -> 7 + self.flatten = nn.Flatten() + self.fc1 = nn.Linear(64 * 7 * 7, 128) + self.relu3 = nn.ReLU() + self.fc2 = nn.Linear(128, num_classes) + + def forward(self, x): + x = self.pool1(self.relu1(self.bn1(self.conv1(x)))) + x = self.pool2(self.relu2(self.bn2(self.conv2(x)))) + # Logits, not softmax: the watched CrossEntropyLoss applies its own. + return self.fc2(self.relu3(self.fc1(self.flatten(x)))) + + +# ============================================================================= +# Train / test +# ============================================================================= +def train(loader, model, optimizer, criterion, device): + """One training step. The projection rides on the criterion call: the + forward hook captured the features, the criterion supplies the ids.""" + with wl.guard_training_context: + inputs, ids, labels = next(loader) + inputs, labels = inputs.to(device), labels.to(device) + optimizer.zero_grad() + logits = model(inputs) + preds = logits.argmax(dim=1, keepdim=True) + loss = criterion(logits.float(), labels.long(), batch_ids=ids, preds=preds) + loss.mean().backward() + optimizer.step() + return loss.mean().detach().cpu().item() + + +def test(loader, model, criterion, metric, device): + """Full pass over the test split. Test images are PLACED in the cloud by the + current encoder but never fitted on, so the layout cannot learn from data + the model is only scored on.""" + total, batches = 0.0, 0 + metric.reset() + for inputs, ids, labels in loader: + with wl.guard_testing_context, torch.no_grad(): + inputs, labels = inputs.to(device), labels.to(device) + logits = model(inputs) + preds = logits.argmax(dim=1, keepdim=True) + total += criterion(logits, labels, batch_ids=ids, preds=preds).mean().item() + metric.update(logits, labels) + batches += 1 + return total / max(1, batches), float(metric.compute() * 100) + + +def projection_status(): + tracker = wl.projection.get_tracker() + if tracker is None: + return "projection off" + stats = tracker.stats() + if stats["disabled_reason"]: + return f"projection stopped: {stats['disabled_reason']}" + return f"umap fits={stats['fits']} placed={stats['samples_written']}" + + +# ============================================================================= +# The stages' building blocks +# ============================================================================= +def sample_table(origin=TRAIN_SPLIT): + """The per-sample table of one split: one row per sample, ``sample_id`` as + the index, with the signals (``signals//...``), ``target``, ``discarded``.""" + wl.drain_signals() # flush pending writes + df = ledgers.get_dataframe().get_df_view() + if "annotation_id" in (df.index.names or []): + df = df.xs(0, level="annotation_id") # the sample-level rows + if "origin" in df.columns: + df = df[df["origin"] == origin] + # The ledger hands ids back as strings and targets as objects; this example's + # ids and classes are numbers, and everything below indexes by them. + df = df.copy() + df.index = pd.Index([int(i) for i in df.index], name="sample_id") + df["target"] = pd.to_numeric(df["target"], errors="coerce") + return df + + +def lowest_loss_ids(table, fraction, column=LOSS_COLUMN): + """Sample ids of the *fraction* of rows with the lowest loss. + + Rows with no loss yet are never picked: an unseen sample is not an easy one. + Returns plain ints, ready for ``wl.discard_samples``. + """ + if column not in table.columns or fraction <= 0: + return [] + loss = pd.to_numeric(table[column], errors="coerce").dropna() + count = int(len(loss) * min(float(fraction), 1.0)) + return [int(i) for i in loss.nsmallest(count).index] if count else [] + + +def save_checkpoint(): + """Write a checkpoint of the model as it is now; returns its step.""" + manager = ledgers.get_checkpoint_manager() + if manager.save_model_checkpoint(force_dump_pending=True, update_manifest=True) is None: + raise RuntimeError("could not write a checkpoint (has training started?)") + return int(ledgers.get_model().get_age()) + + +def reload_checkpoint(step): + """Roll the experiment back to the checkpoint written at *step*: weights, + optimizer, the UMAP encoder, and every sample's loss / seen-count as they + were then. Returns the model's age afterwards.""" + manager = ledgers.get_checkpoint_manager() + if not manager.load_state(manager.get_current_experiment_hash(), target_step=step): + raise RuntimeError(f"could not reload the checkpoint of step {step}") + return int(ledgers.get_model().get_age()) + + +def pca_coordinates(features, dim=2): + """PCA by SVD, in plain torch: ``(N, F)`` features -> ``(N, dim)`` + coordinates, and the share of the variance those ``dim`` axes keep.""" + x = torch.as_tensor(np.asarray(features), dtype=torch.float32) + x = x - x.mean(0) + _, singular, axes = torch.linalg.svd(x, full_matrices=False) + variance = singular ** 2 + return (x @ axes[:dim].T).numpy(), float(variance[:dim].sum() / variance.sum()) + + +def project_with_pca(model, loader, prefix, dim=2): + """Store a PCA of the model's representation as the projection *prefix*. + + WeightsLab runs the loader through the model and collects the features; the + algorithm is the function below -- anything returning ``(N, 2|3)`` works. + Returns ``(stats, features, variance_kept)``; ``features[n]`` belongs to the + sample ``stats["sample_ids"][n]``. + """ + kept = {} + + def pca(features): + coords, kept["variance"] = pca_coordinates(features, dim) + kept["features"] = features + return coords + + stats = wl.project_dataset(model, loader, method=pca, prefix=prefix, verbose=False) + return stats, kept["features"], kept["variance"] + + +def knn_purity(coords, labels, k=10): + """Mean share of each point's *k* nearest neighbours (in ``coords``) that + have its class: 1.0 = every neighbourhood is one class, ~0.1 = ten classes + thrown together. Reads how well a picture keeps the classes apart.""" + points = torch.as_tensor(np.asarray(coords), dtype=torch.float32) + labels = torch.as_tensor(np.asarray(labels)) + k = min(int(k), len(points) - 1) + if k < 1: + return float("nan") + same = 0.0 + for start in range(0, len(points), 2048): # rows at a time: N x N never exists + block = torch.cdist(points[start:start + 2048], points) + block[torch.arange(block.shape[0]), torch.arange(start, start + block.shape[0])] = float("inf") + neighbours = block.topk(k, largest=False).indices + same += (labels[neighbours] == labels[start:start + block.shape[0], None]).float().sum().item() + return same / (len(points) * k) + + +def layout(table, prefix): + """``(ids, coords)`` of the kept samples *prefix* has placed (NaN rows and + discarded samples left out).""" + columns = [f"signals//{prefix}_{axis}" for axis in "xyz" + if f"signals//{prefix}_{axis}" in table.columns] + if not columns: + return [], np.empty((0, 0)) + rows = table[~table["discarded"].astype(bool)] + coords = rows[columns].apply(pd.to_numeric, errors="coerce") + placed = coords.notna().all(axis=1) + return list(rows.index[placed]), coords[placed].to_numpy(dtype=np.float32) + + +def compare(table, prefixes, k=10): + """k-NN class purity of each projection in *prefixes*, over the kept samples + that ALL of them place -- the same points in every picture.""" + layouts = {p: layout(table, p) for p in prefixes} + common = set.intersection(*(set(ids) for ids, _ in layouts.values())) + targets = table["target"] + scores = {} + for prefix, (ids, coords) in layouts.items(): + keep = [n for n, i in enumerate(ids) if i in common] + scores[prefix] = knn_purity(coords[keep], targets.loc[[ids[n] for n in keep]].to_numpy(), k) + return scores, len(common) + + +def plot_side_by_side(table, prefixes): + """A figure with the projections of *prefixes* next to each other, coloured + by class (3-D ones as 3-D axes). The caller shows or saves it.""" + import matplotlib.pyplot as plt + fig = plt.figure(figsize=(6 * len(prefixes), 5.5)) + for slot, prefix in enumerate(prefixes, 1): + ids, coords = layout(table, prefix) + three_d = coords.shape[1] == 3 + ax = fig.add_subplot(1, len(prefixes), slot, projection="3d" if three_d else None) + classes = table["target"].loc[ids].to_numpy() + for c, name in enumerate(CLASS_NAMES): + m = classes == c + ax.scatter(*(coords[m, a] for a in range(coords.shape[1])), s=2, label=name) + ax.set_title(prefix) + ax.set_xticks([]); ax.set_yticks([]) + if three_d: + ax.set_zticks([]) + if slot == 1: + ax.legend(markerscale=5, fontsize=7, loc="best") + fig.tight_layout() + return fig + + +def save_side_by_side(table, prefixes, path): + """:func:`plot_side_by_side`, written to *path* without opening a window.""" + import matplotlib + matplotlib.use("Agg") + import matplotlib.pyplot as plt + fig = plot_side_by_side(table, prefixes) + fig.savefig(path, dpi=130) + plt.close(fig) + return path + + +# ============================================================================= +# The run +# ============================================================================= +def load_parameters(config_path=None): + config_path = config_path or os.path.join(os.path.dirname(os.path.abspath(__file__)), + "config.yaml") + with open(config_path, "r") as fh: + return yaml.safe_load(fh) or {} + + +def run(parameters, steps=None, wait=False, keep_serving=False): + """Run the staged exploration described at the top of this file. + + *parameters* is the config dict (``config.yaml``); *steps* overrides + ``(steps_initial, steps_retrain)``. Returns a summary dict -- checkpoint + step, how many samples were discarded, the purity of each projection. + """ + explore = {"steps_initial": 1000, "steps_retrain": 1000, "discard_fraction": 0.3, + "knn_k": 10} + explore.update(parameters.get("exploration") or {}) + if steps: + explore["steps_initial"], explore["steps_retrain"] = steps + parameters.setdefault("experiment_name", "fashion_mnist_umap_vs_pca") + parameters.setdefault("device", "auto") + parameters["training_steps_to_do"] = None + + # `weightslab start ` exports WEIGHTSLAB_ROOT_LOG_DIR -- honor it, so + # this run lands in the directory the dashboard is watching. + if not parameters.get("root_log_dir"): + parameters["root_log_dir"] = os.environ.get("WEIGHTSLAB_ROOT_LOG_DIR") or tempfile.mkdtemp() + os.makedirs(parameters["root_log_dir"], exist_ok=True) + log_dir = parameters["root_log_dir"] + wl.watch_or_edit(parameters, flag="hyperparameters", poll_interval=1.0) + + if parameters.get("device", "auto") == "auto": + parameters["device"] = torch.device("cuda" if torch.cuda.is_available() else "cpu") + device = torch.device(parameters["device"]) + eval_every = parameters.get("eval_full_to_train_steps_ratio", 250) + enable_h5 = parameters.get("ledger_enable_h5_persistence", True) + + # ---- MODEL: this is where the live UMAP attaches ------------------------ + model = wl.watch_or_edit( + FashionCNN(num_classes=len(CLASS_NAMES)).to(device), flag="model", device=device, + projection=parameters.get("projection") or {}) + lr = parameters.get("optimizer", {}).get("lr", 0.001) + optimizer = wl.watch_or_edit(optim.Adam(model.parameters(), lr=lr), flag="optimizer") + + # ---- DATA --------------------------------------------------------------- + data_root = parameters.get("data_root") or os.path.join(log_dir, "data") + os.makedirs(data_root, exist_ok=True) + train_cfg = parameters.get("data", {}).get("train_loader", {}) + test_cfg = parameters.get("data", {}).get("test_loader", {}) + to_tensor = transforms.ToTensor() + train_dataset = FashionMNISTDataset(root=data_root, train=True, transform=to_tensor, + max_samples=train_cfg.get("max_samples"), id_base=0) + test_dataset = FashionMNISTDataset(root=data_root, train=False, transform=to_tensor, + max_samples=test_cfg.get("max_samples"), + id_base=1_000_000) + train_loader = wl.watch_or_edit( + train_dataset, flag="data", loader_name=TRAIN_SPLIT, + batch_size=train_cfg.get("batch_size", 64), shuffle=train_cfg.get("shuffle", True), + is_training=True, compute_hash=False, preload_labels=True, preload_metadata=True, + enable_h5_persistence=enable_h5) + test_loader = wl.watch_or_edit( + test_dataset, flag="data", loader_name="test_loader", + batch_size=test_cfg.get("batch_size", 256), shuffle=test_cfg.get("shuffle", False), + is_training=False, compute_hash=False, preload_labels=True, preload_metadata=True, + enable_h5_persistence=enable_h5) + + # ---- LOSS / METRIC ------------------------------------------------------ + train_criterion = wl.watch_or_edit( + nn.CrossEntropyLoss(reduction="none"), + flag="loss", signal_name="train-loss-CE", log=True, per_sample=True) + test_criterion = wl.watch_or_edit( + nn.CrossEntropyLoss(reduction="none"), + flag="loss", signal_name="test-loss-CE", log=True, per_sample=True) + metric = wl.watch_or_edit( + Accuracy(task="multiclass", num_classes=len(CLASS_NAMES)).to(device), + flag="metric", signal_name="metric-ACC", log=True) + + wl.serve(serving_grpc=parameters.get("serving_grpc", True)) + wl.start_training(timeout=parameters.get("start_timeout", 3)) + + def train_to(age): + """Train until the model's age reaches *age*; evaluate every so often.""" + bar = tqdm.tqdm(total=max(0, age - model.get_age()), desc="Training", ncols=120, + disable=not parameters.get("tqdm_display", True)) + while model.get_age() < age: + loss = train(train_loader, model, optimizer, train_criterion, device) + bar.update(1) + if model.get_age() % eval_every == 0: + _, accuracy = test(test_loader, model, test_criterion, metric, device) + bar.set_postfix_str(f"loss={loss:.3f} test_acc={accuracy:.1f}% {projection_status()}") + bar.close() + + def banner(text): + print(f"\n[{text}]") + + summary = {"log_dir": log_dir} + print("=" * 72) + print(" FASHION-MNIST: LIVE UMAP vs YOUR OWN PCA") + print(f" train={len(train_dataset)} test={len(test_dataset)} device={device}") + print(f" logs -> {log_dir}") + print(" Studio: weightslab start, then Data Exploration -> Projection.") + print("=" * 72) + + # 1 -- train with the live UMAP --------------------------------------------- + banner(f"1. Train {explore['steps_initial']} steps with the live parametric UMAP") + train_to(explore["steps_initial"]) + summary["checkpoint_step"] = save_checkpoint() + print(f" checkpoint written at step {summary['checkpoint_step']}; {projection_status()}") + + # 2 -- reload those weights ------------------------------------------------- + banner(f"2. Reload the step-{summary['checkpoint_step']} checkpoint") + summary["reloaded_age"] = reload_checkpoint(summary["checkpoint_step"]) + print(f" model age after reload: {summary['reloaded_age']}") + + # 3 -- discard what teaches little ----------------------------------------- + banner(f"3. Discard the {explore['discard_fraction']:.0%} lowest-loss training samples") + table = sample_table() + easy = lowest_loss_ids(table, explore["discard_fraction"]) + if easy: + wl.discard_samples(easy, discarded=True) + summary["discarded"] = len(easy) + print(f" discarded {len(easy)} of {len(table)}") + per_class = table.loc[table.index.isin(easy), "target"].astype(int).value_counts().sort_index() + print(" by class: " + ", ".join(f"{CLASS_NAMES[c]} {n}" for c, n in per_class.items())) + + # 4 -- your own PCA ---------------------------------------------------------- + pca_prefix = f"pca2d_step{summary['reloaded_age']}" + banner(f"4. PCA (2-D) of what is left -> '{pca_prefix}'") + stats, features, variance = project_with_pca(model, train_loader, pca_prefix, dim=2) + summary["pca_samples"] = stats["samples"] + print(f" {stats['samples']} samples x {stats['feature_dim']} features; " + f"2 axes keep {variance:.0%} of the variance") + save_checkpoint() # the PCA is kept with this checkpoint + + # 5 -- compare ---------------------------------------------------------------- + banner("5. Compare the UMAP and the PCA") + table = sample_table() + scores, shared = compare(table, ["umap", pca_prefix], explore["knn_k"]) + summary["purity"] = scores + summary["purity_features"] = knn_purity( + features, table["target"].loc[[int(i) for i in stats["sample_ids"]]].to_numpy(), + explore["knn_k"]) + print(f" neighbours sharing a point's class (k={explore['knn_k']}), over {shared} kept samples:") + print(f" parametric UMAP (3-D) : {scores['umap']:.3f}") + print(f" your PCA (2-D) : {scores[pca_prefix]:.3f}") + print(f" the 128-D features : {summary['purity_features']:.3f} (what both pictures summarize)") + figure = save_side_by_side(table, ["umap", pca_prefix], os.path.join(log_dir, "umap_vs_pca.png")) + print(f" side-by-side picture -> {figure}") + print(" In the studio: Data Exploration -> Projection, pick 'umap' then the PCA.") + if wait: + input(" Look at both in the studio, then press Enter to retrain... ") + + # 6 -- retrain on what is left --------------------------------------------------- + if explore["steps_retrain"] > 0: + target = summary["reloaded_age"] + explore["steps_retrain"] + banner(f"6. Retrain to step {target} on the remaining samples") + train_to(target) + again = f"pca2d_step{model.get_age()}" + project_with_pca(model, train_loader, again, dim=2) + scores, shared = compare(sample_table(), ["umap", again], explore["knn_k"]) + summary["purity_after"] = scores + summary["final_age"] = int(model.get_age()) + print(f" step {model.get_age()}: UMAP {scores['umap']:.3f} | fresh PCA {scores[again]:.3f}" + f" (the live UMAP followed the model; '{pca_prefix}' stayed a snapshot)") + + print("\nDone.") + if keep_serving: + wl.keep_serving() + return summary + + +def main(argv=None): + parser = argparse.ArgumentParser(description=__doc__.split("\n")[0]) + parser.add_argument("--steps", nargs=2, type=int, metavar=("INITIAL", "RETRAIN"), + help="training steps before the discard and after it " + "(default: the config's exploration section, 1000 and 1000)") + parser.add_argument("--wait", action="store_true", + help="stop for Enter after the comparison, to look in the studio") + parser.add_argument("--keep-serving", action="store_true", + help="stay up at the end so you can keep using the studio") + args = parser.parse_args(argv) + run(load_parameters(), steps=args.steps, wait=args.wait, keep_serving=args.keep_serving) + + +if __name__ == "__main__": + main() diff --git a/weightslab/examples/Ultralytics/wl-detection/config.yaml b/weightslab/examples/Ultralytics/wl-detection/config.yaml index f469214f..5a7badc6 100644 --- a/weightslab/examples/Ultralytics/wl-detection/config.yaml +++ b/weightslab/examples/Ultralytics/wl-detection/config.yaml @@ -39,7 +39,7 @@ ledger_flush_interval: 60.0 num_classes: 2 image_size: 320 # data_root: .\data\data.yaml # Uncomment and set the path to your data.yaml file. YOLO format. -data_root: C:\Users\GuillaumePELLUET\Documents\Codes\weightslab_kitchen\guillaume_playground\ws-ultralytics_yolo\data\data.yaml + data: train_loader: batch_size: 4 diff --git a/weightslab/examples/Usecases/wl-3d-lidar-detection/config.yaml b/weightslab/examples/Usecases/wl-3d-lidar-detection/config.yaml index 50796401..38bbc485 100644 --- a/weightslab/examples/Usecases/wl-3d-lidar-detection/config.yaml +++ b/weightslab/examples/Usecases/wl-3d-lidar-detection/config.yaml @@ -2,7 +2,7 @@ device: auto experiment_name: lidar3d_detection_usecase training_steps_to_do: null # null = infinite training until manually stopped (behavior set by the user in main script) -# root_log_dir: # Empty to write in tmp directory, or specify a path to store logs and checkpoints +root_log_dir: C:\Users\GUILLA~1\AppData\Local\Temp\tmpxzyj0tu1 # Empty to write in tmp directory, or specify a path to store logs and checkpoints checkpoint_manager: load_config: false # Disable loading previous checkpoint configuration to allow for changes in the current config.yaml file between experiments. Otherwise changes will be replaced. diff --git a/weightslab/examples/Usecases/wl-fashion-mnist-custom-projection/config.yaml b/weightslab/examples/Usecases/wl-fashion-mnist-custom-projection/config.yaml new file mode 100644 index 00000000..d82972d7 --- /dev/null +++ b/weightslab/examples/Usecases/wl-fashion-mnist-custom-projection/config.yaml @@ -0,0 +1,40 @@ +# Fashion-MNIST: train to convergence, pause from code, then project with your +# own algorithm (t-SNE) from the Studio notebook. The built-in projection is off. +experiment_name: fashion_mnist_custom_projection +device: auto +training_steps_to_do: null +# root_log_dir: # defaults to $WEIGHTSLAB_ROOT_LOG_DIR, else a temp dir +auto_start_training: false # true = start without pressing Play in the studio + +optimizer: + lr: 0.001 + +# When the script pauses itself. +convergence: + eval_every_n_steps: 250 # evaluate the test split this often + patience: 3 # ... and pause after this many evals ... + min_delta: 0.25 # ... without +0.25 accuracy points + pause_after_steps: 6000 # pause by this step at the latest (demo safety net) + +experiment_dump_to_train_steps_ratio: 1000 +skip_checkpoint_load: false +tqdm_display: true +is_training: false + +ledger_enable_flushing_threads: true +ledger_enable_h5_persistence: true +ledger_flush_max_rows: 15000 +ledger_flush_interval: 30.0 + +serving_grpc: true + +data: + train_loader: + shuffle: true + batch_size: 64 + max_samples: 20000 + test_loader: + shuffle: false + batch_size: 256 + max_samples: 2000 # t-SNE is O(N^2): a few thousand points is quick + drop_last: false diff --git a/weightslab/examples/Usecases/wl-fashion-mnist-custom-projection/main.py b/weightslab/examples/Usecases/wl-fashion-mnist-custom-projection/main.py new file mode 100644 index 00000000..1114a989 --- /dev/null +++ b/weightslab/examples/Usecases/wl-fashion-mnist-custom-projection/main.py @@ -0,0 +1,293 @@ +"""Train to convergence, pause from code, then project with YOUR OWN algorithm. + +The flow this example shows: + +1. Fashion-MNIST trains with WeightsLab's built-in projection turned OFF + (``projection=False``) -- the projection here is going to be yours. +2. The script watches test accuracy and, once it stops improving, pauses + training itself with ``wl.pause_training()``. Nothing is lost: the process, + the model and the Studio all stay up; only the next step waits for Play. +3. In Weights Studio's notebook (the button left of the logo), the notebook + ``own-projection`` is ready: it collects the converged model's features, + runs a t-SNE written from scratch in plain torch -- not WeightsLab's -- and + hands the coordinates over with ``wl.save_projection_coords``. +4. Data Exploration -> Projection shows your t-SNE (and the PCA of the last + cell) in the picker: lasso a cluster, and the grid loads those images. + +Running a notebook cell pauses training first, so a cell always sees one state +of the model even if you run it before convergence. + +Run:: + + python main.py # then open the studio and press Play +""" + +import itertools +import logging +import os +import ssl +import tempfile + +try: + ssl.create_default_context() +except ssl.SSLError: + ssl._create_default_https_context = ssl._create_unverified_context + +import torch +import torch.nn as nn +import torch.optim as optim +import tqdm +import yaml +from torch.utils.data import Dataset +from torchmetrics.classification import Accuracy +from torchvision import datasets, transforms + +import weightslab as wl +from projection_notebook import write_projection_notebook + +logging.basicConfig(level=logging.ERROR) +logger = logging.getLogger(__name__) + +CLASS_NAMES = ( + "T-shirt/top", "Trouser", "Pullover", "Dress", "Coat", + "Sandal", "Shirt", "Sneaker", "Bag", "Ankle boot", +) + + +# ============================================================================= +# Dataset / model (same as wl-fashion-mnist-umap) +# ============================================================================= +class FashionMNISTDataset(Dataset): + """Fashion-MNIST yielding ``(image, sample_id, label)``; ids offset per split.""" + + def __init__(self, root, train=True, download=True, transform=None, + max_samples=None, id_base=0): + self.data = datasets.FashionMNIST(root=root, train=train, + download=download, transform=None) + self.transform = transform + self.train = train + self.max_samples = max_samples + self.id_base = id_base + + def __len__(self): + if self.max_samples is not None: + return min(len(self.data), self.max_samples) + return len(self.data) + + def __getitem__(self, idx): + image, label = self.data[idx] + if self.transform: + image = self.transform(image) + return image, self.id_base + idx, label + + def fast_get_label(self, idx): + return int(self.data.targets[idx]) + + def get_metadata(self, idx): + label = int(self.data.targets[idx]) + return {"class_name": CLASS_NAMES[label], + "split": "train" if self.train else "test"} + + +class FashionCNN(nn.Module): + """Two conv blocks, a 128-D representation (``embed``), the classifier.""" + + def __init__(self, num_classes=10): + super().__init__() + self.input_shape = (1, 1, 28, 28) + self.conv1 = nn.Conv2d(1, 32, 3, padding=1) + self.bn1 = nn.BatchNorm2d(32) + self.relu1 = nn.ReLU() + self.pool1 = nn.MaxPool2d(2) + self.conv2 = nn.Conv2d(32, 64, 3, padding=1) + self.bn2 = nn.BatchNorm2d(64) + self.relu2 = nn.ReLU() + self.pool2 = nn.MaxPool2d(2) + self.flatten = nn.Flatten() + self.fc1 = nn.Linear(64 * 7 * 7, 128) + self.relu3 = nn.ReLU() + self.fc2 = nn.Linear(128, num_classes) + + def embed(self, x): + """The representation the notebook projects: fc1 after its ReLU.""" + x = self.pool1(self.relu1(self.bn1(self.conv1(x)))) + x = self.pool2(self.relu2(self.bn2(self.conv2(x)))) + return self.relu3(self.fc1(self.flatten(x))) + + def forward(self, x): + return self.fc2(self.embed(x)) + + +# ============================================================================= +# Convergence +# ============================================================================= +class Plateau: + """Converged = test accuracy has not improved by ``min_delta`` points for + ``patience`` evaluations in a row.""" + + def __init__(self, patience=3, min_delta=0.25): + self.patience = int(patience) + self.min_delta = float(min_delta) + self.best = None + self.stale = 0 + + def update(self, accuracy: float) -> bool: + if self.best is None or accuracy > self.best + self.min_delta: + self.best, self.stale = accuracy, 0 + else: + self.stale += 1 + return self.stale >= self.patience + + +# ============================================================================= +# Train / test +# ============================================================================= +def train(loader, model, optimizer, criterion, device): + with wl.guard_training_context: # blocks here while training is paused + inputs, ids, labels = next(loader) + inputs, labels = inputs.to(device), labels.to(device) + optimizer.zero_grad() + logits = model(inputs) + preds = logits.argmax(dim=1, keepdim=True) + loss = criterion(logits.float(), labels.long(), batch_ids=ids, preds=preds) + loss.mean().backward() + optimizer.step() + return loss.mean().detach().cpu().item() + + +def test(loader, model, criterion, metric, device): + total, batches = 0.0, 0 + metric.reset() + for inputs, ids, labels in loader: + with wl.guard_testing_context, torch.no_grad(): + inputs, labels = inputs.to(device), labels.to(device) + logits = model(inputs) + preds = logits.argmax(dim=1, keepdim=True) + total += criterion(logits, labels, batch_ids=ids, preds=preds).mean().item() + metric.update(logits, labels) + batches += 1 + return total / max(1, batches), float(metric.compute() * 100) + + +def announce_pause(step, accuracy, why, notebook_path): + print("\n" + "=" * 72) + print(f" {why} at step {step} (test accuracy {accuracy:.1f}%).") + print(" Training PAUSED from code (wl.pause_training()). In Weights Studio:") + print(f" 1. open the notebook -> '{notebook_path.stem}' is ready;") + print(" 2. run its cells top to bottom: your own t-SNE of this model;") + print(" 3. Data Exploration -> Projection, pick 'tsne' (or 'pca').") + print(" Press Play in the header to resume training afterwards.") + print("=" * 72 + "\n") + + +# ============================================================================= +# Main +# ============================================================================= +if __name__ == "__main__": + parameters = {} + config_path = os.path.join(os.path.dirname(__file__), "config.yaml") + if os.path.exists(config_path): + with open(config_path, "r") as fh: + parameters = yaml.safe_load(fh) or {} + parameters.setdefault("experiment_name", "fashion_mnist_custom_projection") + parameters.setdefault("device", "auto") + + if not parameters.get("root_log_dir"): + parameters["root_log_dir"] = os.environ.get("WEIGHTSLAB_ROOT_LOG_DIR") or tempfile.mkdtemp() + os.makedirs(parameters["root_log_dir"], exist_ok=True) + log_dir = parameters["root_log_dir"] + + wl.watch_or_edit(parameters, flag="hyperparameters", poll_interval=1.0) + + if parameters.get("device", "auto") == "auto": + parameters["device"] = torch.device("cuda" if torch.cuda.is_available() else "cpu") + device = parameters["device"] + conv = parameters.get("convergence", {}) + eval_every = int(conv.get("eval_every_n_steps", 250)) + plateau = Plateau(conv.get("patience", 3), conv.get("min_delta", 0.25)) + pause_after = conv.get("pause_after_steps") # safety net for the demo + tqdm_display = parameters.get("tqdm_display", True) + enable_h5 = parameters.get("ledger_enable_h5_persistence", True) + + # The built-in projection is OFF: the projection in this run is yours. + model = wl.watch_or_edit(FashionCNN(num_classes=len(CLASS_NAMES)).to(device), + flag="model", device=device, projection=False) + lr = parameters.get("optimizer", {}).get("lr", 0.001) + optimizer = wl.watch_or_edit(optim.Adam(model.parameters(), lr=lr), flag="optimizer") + + data_root = parameters.get("data_root") or os.path.join(log_dir, "data") + os.makedirs(data_root, exist_ok=True) + train_cfg = parameters.get("data", {}).get("train_loader", {}) + test_cfg = parameters.get("data", {}).get("test_loader", {}) + to_tensor = transforms.Compose([transforms.ToTensor()]) + train_dataset = FashionMNISTDataset(root=data_root, train=True, transform=to_tensor, + max_samples=train_cfg.get("max_samples"), id_base=0) + test_dataset = FashionMNISTDataset(root=data_root, train=False, transform=to_tensor, + max_samples=test_cfg.get("max_samples"), + id_base=1_000_000) + train_loader = wl.watch_or_edit( + train_dataset, flag="data", loader_name="train_loader", + batch_size=train_cfg.get("batch_size", 64), shuffle=train_cfg.get("shuffle", True), + is_training=True, compute_hash=False, preload_labels=True, preload_metadata=True, + enable_h5_persistence=enable_h5) + test_loader = wl.watch_or_edit( + test_dataset, flag="data", loader_name="test_loader", + batch_size=test_cfg.get("batch_size", 256), shuffle=test_cfg.get("shuffle", False), + is_training=False, compute_hash=False, preload_labels=True, preload_metadata=True, + enable_h5_persistence=enable_h5) + + train_criterion = wl.watch_or_edit( + nn.CrossEntropyLoss(reduction="none"), + flag="loss", signal_name="train-loss-CE", log=True, per_sample=True) + test_criterion = wl.watch_or_edit( + nn.CrossEntropyLoss(reduction="none"), + flag="loss", signal_name="test-loss-CE", log=True, per_sample=True) + metric = wl.watch_or_edit( + Accuracy(task="multiclass", num_classes=len(CLASS_NAMES)).to(device), + flag="metric", signal_name="metric-ACC", log=True) + + # Before serving, so the Studio's notebook opens on it. + notebook_path = write_projection_notebook(log_dir) + wl.serve(serving_grpc=parameters.get("serving_grpc", True)) + + print("=" * 72) + print(" FASHION-MNIST -> CONVERGE -> PAUSE -> YOUR OWN PROJECTION") + print(f" train={len(train_dataset)} test={len(test_dataset)} device={device}") + print(f" pauses itself when test accuracy plateaus ({plateau.patience} evals of " + f"{eval_every} steps without +{plateau.min_delta} pt)") + print(f" notebook -> {notebook_path}") + print(f" logs -> {log_dir}") + print(" Studio: press Play to start training.") + print("=" * 72 + "\n") + + if parameters.get("auto_start_training", False): + wl.start_training() + + steps = itertools.count() + if tqdm_display: + steps = tqdm.tqdm(steps, desc="Training", ncols=120, + bar_format="{desc}: {n} steps [{elapsed}, {rate_fmt}] | {postfix}") + paused_once = False + test_acc = None + for _ in steps: + train_loss = train(train_loader, model, optimizer, train_criterion, device) + age = model.get_age() + if age > 0 and age % eval_every == 0: + _, test_acc = test(test_loader, model, test_criterion, metric, device) + if not paused_once: + converged = plateau.update(test_acc) + capped = pause_after is not None and age >= int(pause_after) + if converged or capped: + paused_once = True + announce_pause(age, test_acc, + "Converged" if converged else "Step budget reached", + notebook_path) + # The next train() blocks in its guard until Play. + wl.pause_training() + if tqdm_display: + parts = [f"loss={train_loss:.4f}"] + if test_acc is not None: + parts.append(f"test_acc={test_acc:.1f}%") + steps.set_postfix_str(" | ".join(parts)) + + wl.keep_serving() diff --git a/weightslab/examples/Usecases/wl-fashion-mnist-custom-projection/projection_notebook.py b/weightslab/examples/Usecases/wl-fashion-mnist-custom-projection/projection_notebook.py new file mode 100644 index 00000000..55f977e5 --- /dev/null +++ b/weightslab/examples/Usecases/wl-fashion-mnist-custom-projection/projection_notebook.py @@ -0,0 +1,193 @@ +"""The notebook this example drops into the run directory. + +Weights Studio's notebook runs Python IN the training process, against the live +model and data, so these cells see the converged model itself -- no export, no +reload. Running a cell pauses training first (the run has already paused itself +at convergence, so nothing moves while you work). + +The cells are YOUR code: WeightsLab neither collects features nor fits anything +here. It is handed the finished coordinates with ``wl.save_projection_coords`` +and draws them in the Projection Board. +""" + +import json +from pathlib import Path + +NOTEBOOK_NAME = "own-projection.ipynb" + +INTRO = """\ +# Your own projection of the converged model + +WeightsLab's built-in UMAP is **off** in this run (`projection=False`). Everything +below is your own code: collect the representation, run your own t-SNE, then hand +the coordinates to the Projection Board. + +1. **Features** — the 128-D input of the classifier head, for every test image. +2. **Your t-SNE** — exact t-SNE written from scratch in plain torch. +3. **Plug it in** — `wl.save_projection_coords(...)`, then open + *Data Exploration → Projection* and pick `tsne`. +4. *(optional)* a preview here, and a second algorithm (PCA) to compare. + +Running a cell pauses training first; press **Play** in the header to resume. +""" + +FEATURES = """\ +# 1) The representation to project: what flows INTO the classifier (fc2), +# for every test image, from the live converged model. +import torch +from weightslab.backend.ledgers import get_dataloader, get_model + +SPLIT = "test_loader" # or "train_loader" (more points, slower t-SNE) + +net = wl.projection.resolve_module(get_model()) # the FashionCNN itself +device = next(net.parameters()).device +was_training = net.training +net.eval() +features, ids, labels = [], [], [] +with torch.no_grad(): + for x, batch_ids, y in get_dataloader(SPLIT): + features.append(net.embed(x.to(device)).cpu()) + ids.extend(batch_ids.tolist() if hasattr(batch_ids, "tolist") else list(batch_ids)) + labels.append(torch.as_tensor(y)) +net.train(was_training) + +features = torch.cat(features) +labels = torch.cat(labels).numpy() +print(f"{features.shape[0]} samples x {features.shape[1]} features from {SPLIT}") +""" + +TSNE = """\ +# 2) Your own t-SNE (van der Maaten & Hinton, 2008) -- exact, O(N^2), plain torch. +# Fine for a few thousand points; swap in sklearn / openTSNE / umap-learn the +# same way if you prefer -- only the coordinates matter to WeightsLab. +import math, time + +def tsne(X, dim=2, perplexity=30.0, iters=750, lr=200.0, exaggeration=12.0, seed=0, + device="cuda" if torch.cuda.is_available() else "cpu"): + X = torch.as_tensor(X, dtype=torch.float32, device=device) + X = X - X.mean(0) + if X.shape[1] > 50: # usual t-SNE practice + _, _, V = torch.pca_lowrank(X, q=50, center=False) + X = X @ V[:, :50] + n = X.shape[0] + + def sq_dists(A): # one matmul, no sqrt + s = (A * A).sum(1) + return (s[:, None] + s[None, :] - 2.0 * A @ A.T).clamp_min(0.0) + + # P: Gaussian affinities, one bandwidth per point so every row has the target + # perplexity. Distances are shifted by each row's nearest neighbour first, so + # exp() never underflows -- the shift cancels in the normalisation. + D = sq_dists(X) + D.fill_diagonal_(float("inf")) + D = D - D.min(1, keepdim=True).values + D.fill_diagonal_(0.0) + target = math.log(perplexity) + beta = 1.0 / D.mean(1, keepdim=True).clamp_min(1e-12) + lo, hi = torch.zeros_like(beta), torch.full_like(beta, float("inf")) + for _ in range(64): + P = torch.exp(-D * beta) + P.fill_diagonal_(0.0) + sum_p = P.sum(1, keepdim=True) + H = torch.log(sum_p) + beta * (D * P).sum(1, keepdim=True) / sum_p + too_flat = H > target # sharpen those rows + lo = torch.where(too_flat, beta, lo) + hi = torch.where(too_flat, hi, beta) + beta = torch.where(torch.isinf(hi), beta * 2, (lo + hi) / 2) + P = P / sum_p + P = ((P + P.T) / (2 * n)).clamp_min(1e-12) + + # Y: Student-t affinities, gradient descent with momentum and gains. + g = torch.Generator().manual_seed(seed) + Y = (torch.randn(n, dim, generator=g) * 1e-4).to(device) + update, gains = torch.zeros_like(Y), torch.ones_like(Y) + for it in range(iters): + num = 1.0 / (1.0 + sq_dists(Y)) + num.fill_diagonal_(0.0) + Q = (num / num.sum()).clamp_min(1e-12) + W = ((exaggeration if it < 250 else 1.0) * P - Q) * num + grad = 4.0 * (W.sum(1, keepdim=True) * Y - W @ Y) + gains = torch.where((grad > 0) == (update > 0), gains * 0.8, gains + 0.2).clamp_min(0.01) + update = (0.5 if it < 250 else 0.8) * update - lr * gains * grad + Y = Y + update + Y = Y - Y.mean(0) + if it % 150 == 0 or it == iters - 1: + print(f" iter {it:4d} KL(P||Q) = {float((P * torch.log(P / Q)).sum()):.3f}") + return Y.cpu().numpy() + +t0 = time.time() +tsne_coords = tsne(features.numpy(), dim=2) +print(f"t-SNE of {len(tsne_coords)} points in {time.time() - t0:.1f}s") +""" + +PLUG = """\ +# 3) Plug it into the board. Any name works as long as it is not "umap" while +# the built-in projection runs; a 2-D result is drawn on the z = 0 plane. +columns = wl.save_projection_coords(tsne_coords, batch_ids=ids, prefix="tsne") +print("written:", columns) +print("Now open Data Exploration -> Projection: 'tsne' is in the picker.") +print("Lasso a region there to load those images into the grid.") +""" + +PREVIEW = """\ +# 4a) Optional: a quick look here, coloured by class. +import matplotlib.pyplot as plt + +CLASSES = ["T-shirt/top", "Trouser", "Pullover", "Dress", "Coat", + "Sandal", "Shirt", "Sneaker", "Bag", "Ankle boot"] +fig, ax = plt.subplots(figsize=(7, 6)) +for c, name in enumerate(CLASSES): + m = labels == c + ax.scatter(tsne_coords[m, 0], tsne_coords[m, 1], s=3, label=name) +ax.legend(markerscale=4, fontsize=7, loc="best") +ax.set_title("Your t-SNE of the converged representation") +ax.set_xticks([]); ax.set_yticks([]) +""" + +PCA = """\ +# 4b) Optional: a second algorithm, to compare in the board's picker. A 3-D PCA, +# three lines of torch -- linear, so it shows how much of the class structure +# is already linearly separable in the representation. +centered = features - features.mean(0) +_, _, V = torch.pca_lowrank(centered, q=3, center=False) +pca_coords = (centered @ V[:, :3]).numpy() +print("written:", wl.save_projection_coords(pca_coords, batch_ids=ids, prefix="pca")) +""" + + +def _cell(kind, source): + cell = {"cell_type": kind, "metadata": {}, "source": source} + if kind == "code": + cell.update({"execution_count": None, "outputs": []}) + return cell + + +def notebook() -> dict: + return { + "cells": [ + _cell("markdown", INTRO), + _cell("code", FEATURES), + _cell("code", TSNE), + _cell("code", PLUG), + _cell("code", PREVIEW), + _cell("code", PCA), + ], + "metadata": {"kernelspec": {"name": "weightslab", + "display_name": "WeightsLab (shared)"}}, + "nbformat": 4, + "nbformat_minor": 5, + } + + +def write_projection_notebook(root_log_dir) -> Path: + """Write the notebook into *root_log_dir*, unless one of that name is there. + + An existing file is left alone: it may hold your own edits from an earlier + run in the same directory. The Studio opens the most recently written + notebook of the run directory, so this one is what it shows first. + """ + path = Path(root_log_dir) / NOTEBOOK_NAME + if not path.exists(): + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text(json.dumps(notebook(), indent=1), encoding="utf-8") + return path diff --git a/weightslab/examples/Usecases/wl-fashion-mnist-umap/config.yaml b/weightslab/examples/Usecases/wl-fashion-mnist-umap/config.yaml new file mode 100644 index 00000000..6bd8d706 --- /dev/null +++ b/weightslab/examples/Usecases/wl-fashion-mnist-umap/config.yaml @@ -0,0 +1,42 @@ +# Fashion-MNIST classification with the built-in live UMAP projection (beta). +experiment_name: fashion_mnist_live_umap +device: auto +training_steps_to_do: null # no ceiling: training runs until you stop it (studio, CLI, Ctrl+C) +# root_log_dir: # defaults to $WEIGHTSLAB_ROOT_LOG_DIR, else a temp dir + +optimizer: + lr: 0.001 + +# The projection is ON by default; these keys only tune it. Each one overrides +# its WEIGHTSLAB_PROJECTION_* environment variable. See docs/projection.rst. +projection: + every_n_steps: 25 # fit the UMAP encoder every 25 training steps (default 50) + graph_size: 1024 # samples in the kNN graph (default 512) + n_neighbors: 15 + +# Experiment cadence +eval_full_to_train_steps_ratio: 300 +experiment_dump_to_train_steps_ratio: 1000 +skip_checkpoint_load: false +tqdm_display: true +is_training: false + +# Global dataframe storage +ledger_enable_flushing_threads: true +ledger_enable_h5_persistence: true +ledger_flush_max_rows: 15000 +ledger_flush_interval: 30.0 + +serving_grpc: true + +# data_root: ./data # auto-downloaded on first run +data: + train_loader: + shuffle: true + batch_size: 64 + max_samples: 20000 # raise (or drop) for the full 60k + test_loader: + shuffle: false + batch_size: 256 + max_samples: 2000 + drop_last: false diff --git a/weightslab/examples/Usecases/wl-fashion-mnist-umap/main.py b/weightslab/examples/Usecases/wl-fashion-mnist-umap/main.py new file mode 100644 index 00000000..d391fd27 --- /dev/null +++ b/weightslab/examples/Usecases/wl-fashion-mnist-umap/main.py @@ -0,0 +1,301 @@ +"""Fashion-MNIST classification with the live 3-D UMAP projection (beta). + +A plain PyTorch training loop -- watched model/optimizer/loaders/loss/metric, +guard contexts -- and the built-in parametric UMAP running alongside it. There +is no projection code here: wrapping the model with +``wl.watch_or_edit(..., flag="model")`` attaches it, because it is on by +default. The ``projection=`` dict below only tunes it (see config.yaml). + +What you get in Weights Studio, while it trains: + + Data Exploration header -> Projection button opens the 3-D board beside + the data grid. One point per image, coloured by class. Early on the cloud + is one tangle; as the representation forms, ten clusters pull apart -- + and the classes that stay mixed (Shirt / T-shirt / Pullover / Coat) are + exactly the ones the model confuses. Lasso a mixed region to load those + images into the grid. + + The projection reads the INPUT of the last Linear layer (``fc2``): the + 128-D representation the classifier decides from, never the logits. + +Columns ``signals//umap_x|y|z`` appear per sample, so the grid can sort and +filter by them too. The encoder is saved with every checkpoint and restored +with it, so a restart continues the same layout. + +To go further -- reload a checkpoint, discard the lowest-loss samples, lay the +rest out with a PCA of your own and compare it with this UMAP -- see the staged +version, ``examples/PyTorch/wl-fashion-mnist-umap``. + +Run:: + + python main.py # then open the studio and press Play + WEIGHTSLAB_PROJECTION=0 python main.py # the same run, projection removed +""" + +import itertools +import logging +import os +import ssl +import tempfile +import time + +# Windows SSL fix: some Windows cert stores contain malformed ASN1 certs that +# crash ssl.create_default_context(). Fall back to unverified only when broken. +try: + ssl.create_default_context() +except ssl.SSLError: + ssl._create_default_https_context = ssl._create_unverified_context + +import torch +import torch.nn as nn +import torch.optim as optim +import tqdm +import yaml +from torch.utils.data import Dataset +from torchmetrics.classification import Accuracy +from torchvision import datasets, transforms + +import weightslab as wl + +logging.basicConfig(level=logging.ERROR) +logger = logging.getLogger(__name__) + +CLASS_NAMES = ( + "T-shirt/top", "Trouser", "Pullover", "Dress", "Coat", + "Sandal", "Shirt", "Sneaker", "Bag", "Ankle boot", +) + + +# ============================================================================= +# Dataset +# ============================================================================= +class FashionMNISTDataset(Dataset): + """Fashion-MNIST yielding ``(image, sample_id, label)``. + + ``sample_id`` is offset per split (``id_base``) so train and test ids never + collide in the shared ledger. + """ + + def __init__(self, root, train=True, download=True, transform=None, + max_samples=None, id_base=0): + self.data = datasets.FashionMNIST(root=root, train=train, + download=download, transform=None) + self.transform = transform + self.train = train + self.max_samples = max_samples + self.id_base = id_base + + def __len__(self): + if self.max_samples is not None: + return min(len(self.data), self.max_samples) + return len(self.data) + + def __getitem__(self, idx): + image, label = self.data[idx] + if self.transform: + image = self.transform(image) + return image, self.id_base + idx, label + + def fast_get_label(self, idx): + """Lets the ledger read labels at init without decoding every image.""" + return int(self.data.targets[idx]) + + def get_metadata(self, idx): + label = int(self.data.targets[idx]) + return {"class_name": CLASS_NAMES[label], + "split": "train" if self.train else "test"} + + +# ============================================================================= +# Model +# ============================================================================= +class FashionCNN(nn.Module): + """Two conv blocks, then a 128-D representation and the classifier. + + ``fc2`` is the last Linear, so the projection hooks ITS input -- the + 128-D output of ``fc1`` after the ReLU. ``embed()`` returns the same + tensor, for anyone who wants it outside the projection. + """ + + def __init__(self, num_classes=10): + super().__init__() + self.input_shape = (1, 1, 28, 28) + self.conv1 = nn.Conv2d(1, 32, 3, padding=1) + self.bn1 = nn.BatchNorm2d(32) + self.relu1 = nn.ReLU() + self.pool1 = nn.MaxPool2d(2) # 28 -> 14 + self.conv2 = nn.Conv2d(32, 64, 3, padding=1) + self.bn2 = nn.BatchNorm2d(64) + self.relu2 = nn.ReLU() + self.pool2 = nn.MaxPool2d(2) # 14 -> 7 + self.flatten = nn.Flatten() + self.fc1 = nn.Linear(64 * 7 * 7, 128) + self.relu3 = nn.ReLU() + self.fc2 = nn.Linear(128, num_classes) + + def embed(self, x): + x = self.pool1(self.relu1(self.bn1(self.conv1(x)))) + x = self.pool2(self.relu2(self.bn2(self.conv2(x)))) + return self.relu3(self.fc1(self.flatten(x))) + + def forward(self, x): + # Logits, not softmax: the watched CrossEntropyLoss applies its own. + return self.fc2(self.embed(x)) + + +# ============================================================================= +# Train / test +# ============================================================================= +def train(loader, model, optimizer, criterion, device): + """One training step. The projection rides on the criterion call: the + forward hook captured the features, the criterion supplies the ids.""" + with wl.guard_training_context: + inputs, ids, labels = next(loader) + inputs, labels = inputs.to(device), labels.to(device) + optimizer.zero_grad() + logits = model(inputs) + preds = logits.argmax(dim=1, keepdim=True) + loss = criterion(logits.float(), labels.long(), batch_ids=ids, preds=preds) + loss.mean().backward() + optimizer.step() + return loss.mean().detach().cpu().item() + + +def test(loader, model, criterion, metric, device): + """Full pass over the test split. Test images are PLACED in the cloud by + the current encoder but never fitted on, so the layout cannot learn from + data the model is only scored on.""" + total, batches = 0.0, 0 + metric.reset() + for inputs, ids, labels in loader: + with wl.guard_testing_context, torch.no_grad(): + inputs, labels = inputs.to(device), labels.to(device) + logits = model(inputs) + preds = logits.argmax(dim=1, keepdim=True) + total += criterion(logits, labels, batch_ids=ids, preds=preds).mean().item() + metric.update(logits, labels) + batches += 1 + return total / max(1, batches), float(metric.compute() * 100) + + +def projection_status(): + tracker = wl.projection.get_tracker() + if tracker is None: + return "projection off" + stats = tracker.stats() + if stats["disabled_reason"]: + return f"projection stopped: {stats['disabled_reason']}" + return f"umap fits={stats['fits']} placed={stats['samples_written']}" + + +# ============================================================================= +# Main +# ============================================================================= +if __name__ == "__main__": + start_time = time.time() + + parameters = {} + config_path = os.path.join(os.path.dirname(__file__), "config.yaml") + if os.path.exists(config_path): + with open(config_path, "r") as fh: + parameters = yaml.safe_load(fh) or {} + parameters.setdefault("experiment_name", "fashion_mnist_live_umap") + parameters.setdefault("device", "auto") + parameters.setdefault("eval_full_to_train_steps_ratio", 300) + + # `weightslab start ` exports WEIGHTSLAB_ROOT_LOG_DIR -- honor it, so + # this run lands in the directory the dashboard is watching. + if not parameters.get("root_log_dir"): + parameters["root_log_dir"] = os.environ.get("WEIGHTSLAB_ROOT_LOG_DIR") or tempfile.mkdtemp() + os.makedirs(parameters["root_log_dir"], exist_ok=True) + log_dir = parameters["root_log_dir"] + + wl.watch_or_edit(parameters, flag="hyperparameters", poll_interval=1.0) + + if parameters.get("device", "auto") == "auto": + parameters["device"] = torch.device("cuda" if torch.cuda.is_available() else "cpu") + device = parameters["device"] + eval_ratio = parameters.get("eval_full_to_train_steps_ratio", 300) + tqdm_display = parameters.get("tqdm_display", True) + enable_h5 = parameters.get("ledger_enable_h5_persistence", True) + + # ---- MODEL ------------------------------------------------------------- + # This is where the projection attaches. Nothing else is needed; the dict + # only tunes it (every_n_steps, graph_size, ...). `projection=False` would + # turn it off for this model, `projection="fc1"` would hook another layer. + model = wl.watch_or_edit( + FashionCNN(num_classes=len(CLASS_NAMES)).to(device), + flag="model", device=device, + projection=parameters.get("projection") or {}, + ) + lr = parameters.get("optimizer", {}).get("lr", 0.001) + optimizer = wl.watch_or_edit(optim.Adam(model.parameters(), lr=lr), flag="optimizer") + + # ---- DATA -------------------------------------------------------------- + data_root = parameters.get("data_root") or os.path.join(log_dir, "data") + os.makedirs(data_root, exist_ok=True) + train_cfg = parameters.get("data", {}).get("train_loader", {}) + test_cfg = parameters.get("data", {}).get("test_loader", {}) + to_tensor = transforms.Compose([transforms.ToTensor()]) + + train_dataset = FashionMNISTDataset( + root=data_root, train=True, transform=to_tensor, + max_samples=train_cfg.get("max_samples"), id_base=0) + test_dataset = FashionMNISTDataset( + root=data_root, train=False, transform=to_tensor, + max_samples=test_cfg.get("max_samples"), id_base=1_000_000) + + train_loader = wl.watch_or_edit( + train_dataset, flag="data", loader_name="train_loader", + batch_size=train_cfg.get("batch_size", 64), shuffle=train_cfg.get("shuffle", True), + is_training=True, compute_hash=False, preload_labels=True, preload_metadata=True, + enable_h5_persistence=enable_h5) + test_loader = wl.watch_or_edit( + test_dataset, flag="data", loader_name="test_loader", + batch_size=test_cfg.get("batch_size", 256), shuffle=test_cfg.get("shuffle", False), + is_training=False, compute_hash=False, preload_labels=True, preload_metadata=True, + enable_h5_persistence=enable_h5) + + # ---- LOSS / METRIC ----------------------------------------------------- + train_criterion = wl.watch_or_edit( + nn.CrossEntropyLoss(reduction="none"), + flag="loss", signal_name="train-loss-CE", log=True, per_sample=True) + test_criterion = wl.watch_or_edit( + nn.CrossEntropyLoss(reduction="none"), + flag="loss", signal_name="test-loss-CE", log=True, per_sample=True) + metric = wl.watch_or_edit( + Accuracy(task="multiclass", num_classes=len(CLASS_NAMES)).to(device), + flag="metric", signal_name="metric-ACC", log=True) + + wl.serve(serving_grpc=parameters.get("serving_grpc", True)) + + print("=" * 72) + print(" FASHION-MNIST + LIVE 3-D UMAP PROJECTION (beta)") + print(f" train={len(train_dataset)} test={len(test_dataset)} device={device}") + print(f" {projection_status()} | eval every {eval_ratio} steps") + print(f" logs -> {log_dir}") + print(" Studio: press Play, then Data Exploration -> Projection.") + print("=" * 72 + "\n") + + # Uncomment to start without pressing Play in the studio. + # wl.start_training() + + # Training runs until YOU stop it (studio pause button, CLI, Ctrl+C). + steps = itertools.count() + if tqdm_display: + steps = tqdm.tqdm(steps, desc="Training", ncols=140, + bar_format="{desc}: {n} steps [{elapsed}, {rate_fmt}] | {postfix}") + test_loss = test_acc = None + for _ in steps: + train_loss = train(train_loader, model, optimizer, train_criterion, device) + age = model.get_age() + if age > 0 and age % eval_ratio == 0: + test_loss, test_acc = test(test_loader, model, test_criterion, metric, device) + if tqdm_display: + parts = [f"loss={train_loss:.4f}"] + if test_acc is not None: + parts.append(f"test_acc={test_acc:.1f}%") + parts.append(projection_status()) + steps.set_postfix_str(" | ".join(parts)) + + wl.keep_serving() diff --git a/weightslab/examples/Usecases/wl-mnist-at-scale/main.py b/weightslab/examples/Usecases/wl-mnist-at-scale/main.py new file mode 100644 index 00000000..350f5941 --- /dev/null +++ b/weightslab/examples/Usecases/wl-mnist-at-scale/main.py @@ -0,0 +1,470 @@ +"""MNIST at scale, with no training: a simulated run to browse in Weights Studio. + +Builds a dataset of ``--samples`` samples (1M by default, split 80/10/10 into +train/val/test) out of MNIST's 70,000 images -- each sample is one of them with +its own small, deterministic shift and contrast, so neighbouring samples do not +look identical -- and fills in what a real training run would have left behind, +without training anything: + + * the run itself: ``--steps`` steps (250k) at batch size ``--batch-size`` (4), + i.e. 1M sample visits, 1.25 epochs of the train split. Curves for + ``train-loss`` / ``train-acc`` at every step (each step a real batch of 4 + train samples, so the batch accuracy moves in quarters, as it would), and + ``val-loss`` / ``val-acc`` / ``test-loss`` / ``test-acc`` every + ``--eval-every`` steps over the whole split -- with the val/test loss + turning back up late in the run, a little overfitting to look at; + * per sample: the loss of every training visit (its trajectory), when it was + last seen and how often; the final model's prediction, loss and confidence + (~97% accurate on val/test, making MNIST's real mistakes: 4<->9, 3<->5, + 7<->1, ...); + * a 3-D projection (``signals//umap_{x,y,z}``): one cluster per class, a few + sub-clusters per class (writing styles), misclassified samples pulled toward + the class they were mistaken for; + * two tags: ``hard`` (the highest final losses) and ``label_noise`` + (confident and wrong). + +Every number is derived from one per-sample difficulty, so the curves end +exactly on what the grid and the projection show: the last val-acc point is +the share of val samples the grid shows as correctly predicted. + +Then it serves the Studio and waits. Sample statistics are not written to H5, +so the run starts the same every time. + +Run:: + + python main.py --mnist-root # 1M samples + weightslab start # another terminal: open the Studio + +Memory, measured at 1M samples with the Studio connected (grid, projection, +lasso, selection applied): 2.4 GB committed, about 2.3 GB per million samples +on top of ~0.5 GB of libraries. The ledger keeps several Python objects per +sample and the data board builds its own copy of the table. So 10M samples +needs ~24 GB of free RAM -- a 32 GB machine with little else open -- while a +16 GB laptop holds 2-4M depending on what else is running. Startup warns when +the request will not fit, and prints the resident memory after each stage. +Registration takes ~40 s per million samples. +""" + +import argparse +import os +import ssl +import sys +import tempfile +import time + +try: + ssl.create_default_context() +except ssl.SSLError: + ssl._create_default_https_context = ssl._create_unverified_context + +import numpy as np +import pandas as pd +import torch +from torch.utils.data import Dataset + +# MNIST's classic confusions: for each class, the classes it is mistaken for +# and how often, relative to each other. +CONFUSIONS = { + 0: {6: 3, 2: 1, 8: 1}, 1: {7: 4, 2: 1, 8: 1}, 2: {7: 3, 3: 2, 8: 1}, + 3: {5: 4, 8: 2, 2: 1}, 4: {9: 6, 6: 1, 1: 1}, 5: {3: 4, 6: 2, 8: 1}, + 6: {0: 3, 5: 2, 4: 1}, 7: {1: 3, 9: 3, 2: 2}, 8: {3: 3, 5: 2, 9: 2}, + 9: {4: 6, 7: 2, 8: 1}, +} +STYLES = 3 # sub-clusters per class in the projection +CHANCE_LOSS = float(np.log(10)) +GB_PER_MILLION = 2.3 # measured with the Studio connected (see above) +BASE_GB = 0.5 + + +def rss_gb() -> float: + try: + import psutil + return psutil.Process().memory_info().rss / 2 ** 30 + except Exception: + return float("nan") + + +def stage(name: str, started: float) -> float: + now = time.perf_counter() + print(f"[mnist-at-scale] {name:<46} {now - started:7.1f} s RSS {rss_gb():5.1f} GB", + flush=True) + return now + + +def warn_if_too_big(samples: int) -> None: + """Say so up front when the run will not fit in free memory: past that + point Windows pages the ledger out, and everything slows to a crawl.""" + try: + import psutil + available = psutil.virtual_memory().available / 2 ** 30 + except Exception: + return + needed = BASE_GB + GB_PER_MILLION * samples / 1e6 + if needed > available: + fits = max(0.0, (available - BASE_GB) / GB_PER_MILLION) + print(f"[mnist-at-scale] WARNING: {samples:,} samples need ~{needed:.1f} GB, " + f"{available:.1f} GB is free -- about {fits:.1f}M samples fit right now. " + f"Expect paging (slow) until other applications free memory.", flush=True) + + +def load_mnist(root: str): + """All 70,000 MNIST images (uint8, [70000, 28, 28]) and labels.""" + from torchvision import datasets + images, labels = [], [] + for train in (True, False): + split = datasets.MNIST(root, train=train, download=True) + images.append(split.data.numpy()) + labels.append(split.targets.numpy()) + return np.concatenate(images), np.concatenate(labels).astype(np.int64) + + +def mix(x: np.ndarray, salt: int) -> np.ndarray: + """Deterministic per-sample pseudo-random bits (splitmix64).""" + x = (np.asarray(x).astype(np.uint64) + np.uint64(salt)) * np.uint64(0x9E3779B97F4A7C15) + x ^= x >> np.uint64(30) + x *= np.uint64(0xBF58476D1CE4E5B9) + x ^= x >> np.uint64(27) + return x + + +def unit(bits: np.ndarray) -> np.ndarray: + """Pseudo-random bits -> uniform floats in [0, 1).""" + return (bits >> np.uint64(11)).astype(np.float64) / float(1 << 53) + + +class ScaledMNIST(Dataset): + """``count`` samples, sample ``i`` being MNIST image ``(start + i) % 70000`` + shifted by up to 2 px and rescaled in contrast -- decided by its index, so + the same sample always looks the same. + + ``fast_get_label`` is what lets the ledger register millions of samples + without decoding a single image. + """ + + def __init__(self, images: np.ndarray, labels: np.ndarray, start: int, count: int): + self.images, self.labels = images, labels + self.start, self.count = start, count + + def __len__(self): + return self.count + + def _base(self, idx: int) -> int: + return (self.start + idx) % len(self.images) + + def __getitem__(self, idx): + idx = int(idx) + bits = int(mix(np.array([self.start + idx]), 1)[0]) + dx, dy = bits % 5 - 2, (bits >> 8) % 5 - 2 + contrast = 0.75 + ((bits >> 16) % 64) / 256.0 + image = np.roll(self.images[self._base(idx)], (dy, dx), axis=(0, 1)) + tensor = torch.from_numpy(image.astype(np.float32) * (contrast / 255.0)).unsqueeze(0) + return tensor, int(self.labels[self._base(idx)]) + + def fast_get_label(self, idx): + return None, None, int(self.labels[self._base(int(idx))]) + + +# ============================================================================= +# The model, simulated: one difficulty per sample drives everything +# ============================================================================= +def final_state(global_index: np.ndarray, base: np.ndarray, target: np.ndarray, + seed: int, training: bool) -> dict: + """The trained model's view of these samples: prediction, loss, + confidence, place in the projection -- and the difficulty behind them, + which the training curves are built from too.""" + rng = np.random.default_rng(seed) + + # Difficulty: mostly easy, a long tail. Partly the image itself (some MNIST + # digits are genuinely ambiguous), partly this sample's distortion. The + # model fits its training samples better than it generalises. + image_hardness = unit(mix(base, 7)) ** 6 + distortion = unit(mix(global_index, 11)) ** 3 + difficulty = np.clip(0.75 * image_hardness + 0.35 * distortion, 0, 1) + p_wrong = (0.0006 + 0.08 * difficulty ** 2) if training else (0.002 + 0.22 * difficulty ** 2) + wrong = unit(mix(global_index, 13)) < p_wrong + + prediction = target.copy() + pick = unit(mix(global_index, 17)) + for cls, confused in CONFUSIONS.items(): + rows = np.flatnonzero(wrong & (target == cls)) + if not len(rows): + continue + choices = np.array(list(confused.keys())) + weights = np.cumsum(np.array(list(confused.values()), dtype=float)) + prediction[rows] = choices[np.searchsorted(weights / weights[-1], pick[rows])] + + jitter = unit(mix(global_index, 19)) + confidence = np.where(wrong, 0.45 + 0.5 * jitter, 0.9995 - 0.5 * difficulty ** 2 * jitter) + p_true = np.where(wrong, (1 - confidence) * (0.3 + 0.6 * jitter), confidence) + loss = -np.log(np.clip(p_true, 1e-6, 1.0)) + + # Projection: classes on a sphere, writing styles around each class, + # distortion spreading a sample out, mistakes pulled toward the class they + # were mistaken for. Same seed for every split: one embedding space. + angle = rng.uniform(0, 2 * np.pi, (10, 2)) + centres = 12.0 * np.stack([np.sin(angle[:, 0]) * np.cos(angle[:, 1]), + np.sin(angle[:, 0]) * np.sin(angle[:, 1]), + np.cos(angle[:, 0])], axis=1) + styles = rng.normal(0, 2.2, (10, STYLES, 3)) + style = (mix(base, 23) % np.uint64(STYLES)).astype(np.int64) + noise = np.stack([unit(mix(global_index, s)) for s in (29, 31, 37)], axis=1) + spread = 0.9 + 1.8 * difficulty[:, None] + xyz = centres[target] + styles[target, style] + (noise - 0.5) * 2.0 * spread + pull = np.where(wrong, 0.35 + 0.5 * jitter, 0.0)[:, None] + xyz = xyz + pull * (centres[prediction] - centres[target]) + + return {"difficulty": difficulty, "wrong": wrong, "prediction": prediction, + "loss": loss, "confidence": confidence, "xyz": xyz} + + +def untrained(step: np.ndarray, steps: int) -> np.ndarray: + """How far from trained the model is at *step*: 1 at the start, exactly 0 + at the last step, so the final evaluation IS the final per-sample state. + A fast early drop, then a long slow tail.""" + step = np.asarray(step, dtype=np.float64) + tail = lambda s: 1.0 / (1.0 + s / 15_000.0) + slow = (tail(step) - tail(steps)) / (1.0 - tail(steps)) + return np.clip(0.7 * np.exp(-step / 4_000.0) * (1 - step / steps) + 0.3 * slow, 0.0, 1.0) + + +def overfit(step: np.ndarray, steps: int) -> np.ndarray: + """Held-out loss creeping back up over the last 40% of the run.""" + s = np.asarray(step, dtype=np.float64) + return np.clip((s - 0.6 * steps) / (0.4 * steps), 0.0, 1.0) ** 2 + + +def loss_at(state: dict, g, late) -> np.ndarray: + """Expected per-sample loss when the model is *g* from trained.""" + start = CHANCE_LOSS * (0.7 + 0.6 * state["difficulty"]) + held_out = 0.04 * late * (0.3 + 2.0 * state["difficulty"] + 3.0 * state["wrong"]) + return state["loss"] + g * (start - state["loss"]) + held_out + + +def p_correct_at(state: dict, g) -> np.ndarray: + """Chance the model gets the sample right when it is *g* from trained.""" + return (~state["wrong"]) * (1 - g) + 0.1 * g + + +def simulate_run(train: dict, held_out: dict, steps: int, batch_size: int, + eval_every: int, seed: int) -> dict: + """The curves of the run, and every training visit of every sample.""" + rng = np.random.default_rng(seed + 1) + n_train = len(train["difficulty"]) + visits = steps * batch_size + + # Shuffled epochs over the train split, cut at the last step. + order = np.concatenate([rng.permutation(n_train) + for _ in range(-(-visits // n_train))])[:visits] + step_of = np.arange(visits) // batch_size + 1 + g = untrained(step_of, steps) + expected = loss_at({k: v[order] for k, v in train.items()}, g, 0.0) + visit_loss = expected * rng.lognormal(-0.06, 0.35, visits) + visit_right = rng.random(visits) < p_correct_at({"wrong": train["wrong"][order]}, g) + + per_step_loss = visit_loss.reshape(steps, batch_size) + curves = { + "train-loss": (np.arange(1, steps + 1), per_step_loss.mean(axis=1), + per_step_loss.min(axis=1), per_step_loss.max(axis=1)), + "train-acc": (np.arange(1, steps + 1), + visit_right.reshape(steps, batch_size).mean(axis=1), None, None), + } + eval_steps = np.arange(eval_every, steps + 1, eval_every) + if eval_steps[-1] != steps: + eval_steps = np.append(eval_steps, steps) + for split, state in held_out.items(): + losses, accs = [], [] + for s in eval_steps: + g_s, late = float(untrained(s, steps)), float(overfit(s, steps)) + losses.append(float(loss_at(state, g_s, late).mean())) + accs.append(float(p_correct_at(state, g_s).mean())) + # Evaluation noise, none at the last point: it must read exactly what + # the grid shows for the final model. + jitter = 1 + rng.normal(0, 0.004, len(eval_steps)) + jitter[-1] = 1.0 + curves[f"{split}-loss"] = (eval_steps, np.array(losses) * jitter, None, None) + curves[f"{split}-acc"] = (eval_steps, np.clip(np.array(accs) * jitter, 0, 1), None, None) + return {"curves": curves, "order": order, "step_of": step_of, "visit_loss": visit_loss} + + +def log_curves(curves: dict, sample_rows: tuple, run_seconds: float, batch_size: int) -> str: + """Write the run into the Studio's signal history, as training would have. + + Rows go through the logger's own staging and flush -- the path every + logged point takes -- but not one ``add_scalars`` call per point: that also + queues each point for live streaming, half a million entries no browser + is waiting for. + """ + from weightslab.backend.ledgers import get_logger + logger = get_logger() + manager = getattr(logger, "chkpt_manager", None) + exp_hash = manager.get_current_experiment_hash() if manager else None + now = int(time.time()) + total = max(int(c[0][-1]) for c in curves.values()) + with logger._lock: + for name, (steps, values, low, high) in curves.items(): + logger.graph_names.add(name) + stamps = now - (run_seconds * (1 - steps / total)).astype(np.int64) + for i in range(len(steps)): + logger._stage_signal_row( + name, exp_hash, int(steps[i]), float(values[i]), int(stamps[i]), + False, False, "", [], "", + sample_count=batch_size if low is not None else 0, + value_min=None if low is None else float(low[i]), + value_max=None if high is None else float(high[i])) + logger._flush_stage() + name, sample_ids, steps, values = sample_rows + logger.ingest_per_sample(name, exp_hash, zip(sample_ids, steps, values)) + logger._flush_stage() + return exp_hash + + +def parse_args(argv=None): + p = argparse.ArgumentParser(description=__doc__.splitlines()[0]) + p.add_argument("--samples", type=int, default=1_000_000) + p.add_argument("--val-fraction", type=float, default=0.1) + p.add_argument("--test-fraction", type=float, default=0.1) + p.add_argument("--steps", type=int, default=250_000, help="Simulated training steps.") + p.add_argument("--batch-size", type=int, default=4, help="Simulated training batch size.") + p.add_argument("--eval-every", type=int, default=2_500, + help="Steps between two (simulated) val/test evaluations.") + p.add_argument("--mnist-root", default=os.environ.get("WL_DATA_ROOT", "./data"), + help="Directory holding MNIST/raw (downloaded there if missing).") + p.add_argument("--chunk", type=int, default=1_000_000, + help="Samples written to the ledger per batch (bounds peak memory).") + p.add_argument("--hard-share", type=float, default=0.01, + help="Share of samples tagged 'hard' (the highest final losses).") + p.add_argument("--seed", type=int, default=0) + p.add_argument("--log-dir", default=None) + p.add_argument("--grpc-port", type=int, default=50051) + p.add_argument("--no-serve", action="store_true", + help="Build everything, report, and exit (for measuring).") + return p.parse_args(argv) + + +def main(argv=None) -> int: + args = parse_args(argv) + warn_if_too_big(args.samples) + started = time.perf_counter() + log_dir = args.log_dir or tempfile.mkdtemp(prefix="wl-mnist-at-scale-") + os.makedirs(log_dir, exist_ok=True) + os.environ["WEIGHTSLAB_ROOT_LOG_DIR"] = log_dir + + import weightslab as wl + from weightslab.backend.ledgers import get_dataframe + from weightslab.projection.registry import register_prefix + + # The run's settings, shown in the Studio's hyperparameters panel. H5 off: + # a million rows to disk on every flush would be most of the startup, for + # a run that never needs reloading. + hp = {"experiment_name": "mnist-at-scale (simulated)", "root_log_dir": log_dir, + "ledger_enable_h5_persistence": False, "ledger_flush_max_rows": 1_000_000, + "batch_size": args.batch_size, "training_steps": args.steps, + "eval_every": args.eval_every, "optimizer": "adam", "learning_rate": 1e-3, + "architecture": "small CNN (simulated -- nothing is trained)", + # Where the Studio's hyperparameters panel reads each loader's batch size. + "data": {"train_loader": {"batch_size": args.batch_size}, + "val_loader": {"batch_size": 256}, "test_loader": {"batch_size": 256}}} + wl.watch_or_edit(hp, flag="hyperparameters", defaults=hp) + + images, labels = load_mnist(args.mnist_root) + t = stage(f"MNIST loaded ({len(images):,} images)", started) + + n_val = int(round(args.samples * args.val_fraction)) + n_test = int(round(args.samples * args.test_fraction)) + n_train = args.samples - n_val - n_test + splits = [("train_loader", 0, n_train, True), + ("val_loader", n_train, n_val, False), + ("test_loader", n_train + n_val, n_test, False)] + loaders, states = {}, {} + for name, start, count, training in splits: + if count <= 0: + continue + loaders[name] = wl.watch_or_edit( + ScaledMNIST(images, labels, start, count), + flag="data", loader_name=name, + batch_size=args.batch_size if training else 256, shuffle=training, + is_training=training, compute_hash=False, + preload_labels=True, preload_metadata=False, root_log_dir=log_dir, + ) + global_index = np.arange(start, start + count) + base = global_index % len(images) + states[name] = final_state(global_index, base, labels[base], args.seed, training) + t = stage(f"registered {name} ({count:,} samples)", t) + + # --- the run -------------------------------------------------------------- + run = simulate_run(states["train_loader"], + {"val": states["val_loader"], "test": states["test_loader"]}, + args.steps, args.batch_size, args.eval_every, args.seed) + train_ids = np.asarray(loaders["train_loader"].wrapped_dataset.unique_ids) + exp_hash = log_curves( + run["curves"], + ("train-loss", train_ids[run["order"]], run["step_of"], run["visit_loss"]), + run_seconds=args.steps * 0.05, batch_size=args.batch_size) + t = stage(f"run logged ({args.steps:,} steps, batch {args.batch_size})", t) + + # Per training sample: its last visit, how many, and that visit's loss. + order, step_of = run["order"], run["step_of"] + last_seen = np.full(n_train, -1, dtype=np.int64) + last_seen[order] = step_of # later visits overwrite earlier + nb_seen = np.bincount(order, minlength=n_train) + last_loss = np.full(n_train, np.nan, dtype=np.float32) + last_loss[order] = run["visit_loss"] + + # --- per-sample state ------------------------------------------------------- + ledger = get_dataframe() + ledger.register_boolean_tag("label_noise") + ledger.register_boolean_tag("hard") + all_loss = np.concatenate([states[name]["loss"] for name in states]) + hard_cut = float(np.quantile(all_loss, 1 - args.hard_share)) + for name, start, count, training in splits: + if name not in loaders: + continue + state, ids = states[name], np.asarray(loaders[name].wrapped_dataset.unique_ids) + late = overfit(args.steps, args.steps) if not training else 0.0 + final_loss = loss_at(state, 0.0, late) + for lo in range(0, count, args.chunk): + hi = min(count, lo + args.chunk) + part = slice(lo, hi) + columns = { + "prediction": state["prediction"][part], + "signals//loss": final_loss[part].astype(np.float32), + "signals//confidence": state["confidence"][part].astype(np.float32), + "signals//umap_x": state["xyz"][part, 0].astype(np.float32), + "signals//umap_y": state["xyz"][part, 1].astype(np.float32), + "signals//umap_z": state["xyz"][part, 2].astype(np.float32), + "tag:label_noise": state["wrong"][part] & (state["confidence"][part] > 0.9), + "tag:hard": final_loss[part] >= hard_cut, + } + if training: + columns["signals//train-loss"] = last_loss[part] + columns["last_seen"] = last_seen[part] + columns["nb_seen"] = nb_seen[part] + frame = pd.DataFrame(columns, index=pd.MultiIndex.from_arrays( + [ids[lo:hi], np.zeros(hi - lo, dtype=np.int64)], + names=["sample_id", "annotation_id"])) + ledger.upsert_df(frame, origin=name) + del frame, columns + t = stage(f"predictions + projection written ({name})", t) + + # Tell the Projection board "umap" is a projection, not just three columns. + register_prefix("umap", root_log_dir=log_dir) + + curves = run["curves"] + summary = ", ".join(f"{k} {curves[k][1][-1]:.3f}" for k in + ("train-loss", "val-loss", "val-acc", "test-loss", "test-acc")) + total = sum(len(l.wrapped_dataset) for l in loaders.values()) + print(f"[mnist-at-scale] {total:,} samples ready in {time.perf_counter() - started:.0f} s " + f"(RSS {rss_gb():.1f} GB). Run {exp_hash}: {summary}.", flush=True) + print(f"[mnist-at-scale] log dir: {log_dir}", flush=True) + if args.no_serve: + return 0 + + wl.serve(serving_grpc=True, serving_cli=True, grpc_port=args.grpc_port) + print("[mnist-at-scale] serving -- open the Studio with `weightslab start` " + f"(backend port {args.grpc_port}). Ctrl+C to stop.", flush=True) + wl.keep_serving() + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/weightslab/examples/Usecases/wl-parametric-umap/main.py b/weightslab/examples/Usecases/wl-parametric-umap/main.py new file mode 100644 index 00000000..4c535384 --- /dev/null +++ b/weightslab/examples/Usecases/wl-parametric-umap/main.py @@ -0,0 +1,261 @@ +"""Live parametric-UMAP projection, end to end, on synthetic clustered data. + +Nothing here opts the projection in: wrapping the model with +``wl.watch_or_edit(..., flag="model")`` attaches it, because the feature is on +by default. What this example shows is that the coordinates land on the samples +as ordinary per-sample signals, and that the 3-D layout recovers the structure +the data actually has. + +Synthetic blobs rather than MNIST on purpose: the ground-truth cluster of every +sample is known, so the run can *assert* that the projection separated them +instead of asking you to squint at a scatter plot. + +Run it:: + + python main.py + +Turn the projection off and watch the columns disappear:: + + WEIGHTSLAB_PROJECTION=0 python main.py + +Other knobs: ``WEIGHTSLAB_PROJECTION_EVERY`` (fit cadence, default 50 steps), +``WEIGHTSLAB_PROJECTION_DIM`` (3), ``WEIGHTSLAB_PROJECTION_NEIGHBORS`` (15), +``WEIGHTSLAB_PROJECTION_GRAPH`` (512). + +It ends by plugging in projections of its own -- the same path a t-SNE or a +umap-learn layout of yours would take -- so the board's picker offers +``umap``, ``umap_trunk``, ``pca`` and ``pca_inputs`` side by side. +""" + +import os +import shutil +import tempfile + +import numpy as np +import torch +import torch.nn as nn +import torch.optim as optim +from torch.utils.data import Dataset + +# Fit often, because this demo only trains for a few hundred steps. +os.environ.setdefault("WEIGHTSLAB_PROJECTION_EVERY", "5") + +import weightslab as wl # noqa: E402 +from weightslab.proto import experiment_service_pb2 as pb2 # noqa: E402 +from weightslab.trainer.services import projection_service as ps # noqa: E402 + +LOSS = "loss_sample" +CLUSTERS = 6 +PER_CLUSTER = 256 +DIM = 32 +BATCH = 64 +OUT = os.environ.get("WL_UMAP_OUT") or os.path.join(tempfile.gettempdir(), "wl_umap_demo") + + +class Blobs(Dataset): + """Well-separated Gaussian blobs in 32-D. Yields ``(x, uid, label)`` -- the + uid is what ties a row in the dataframe to a point in the projection.""" + + def __init__(self, clusters=CLUSTERS, per_cluster=PER_CLUSTER, dim=DIM, seed=0): + rng = np.random.default_rng(seed) + centers = rng.normal(scale=5.0, size=(clusters, dim)) + self.x = np.concatenate( + [centers[c] + rng.normal(scale=0.6, size=(per_cluster, dim)) + for c in range(clusters)]).astype(np.float32) + self.y = np.concatenate( + [np.full(per_cluster, c) for c in range(clusters)]).astype(np.int64) + + def __len__(self): + return len(self.x) + + def __getitem__(self, i): + return torch.from_numpy(self.x[i]), i, int(self.y[i]) + + def fast_get_label(self, i): + return int(self.y[i]) + + +class Net(nn.Module): + """trunk -> 16-D features -> classifier. The projection auto-hooks the + classifier's INPUT, so it watches the 16-D representation, not the logits.""" + + def __init__(self, dim=DIM, feat=16, classes=CLUSTERS): + super().__init__() + self.input_shape = (1, dim) + self.trunk = nn.Sequential(nn.Linear(dim, 64), nn.ReLU(), nn.Linear(64, feat), nn.ReLU()) + self.classifier = nn.Linear(feat, classes) + + def forward(self, x): + return self.classifier(self.trunk(x)) + + +def main(): + shutil.rmtree(OUT, ignore_errors=True) + os.makedirs(OUT, exist_ok=True) + device = "cuda" if torch.cuda.is_available() else "cpu" + + hp = { + "experiment_name": "parametric-umap-demo", + "device": device, + "root_log_dir": os.path.join(OUT, "wl_logs"), + "serving_grpc": False, + "serving_cli": False, + "ledger_enable_h5_persistence": False, + "experiment_dump_to_train_steps_ratio": 10_000_000, + "data": {"train_loader": {"batch_size": BATCH, "shuffle": True}}, + } + wl.watch_or_edit(hp, flag="hyperparameters", defaults=hp) + + dataset = Blobs() + # Wrapping the model is all it takes -- the projection attaches here. + model = wl.watch_or_edit(Net().to(device), flag="model", device=device) + opt = wl.watch_or_edit(optim.Adam(model.parameters(), lr=3e-3), flag="optimizer") + loader = wl.watch_or_edit(dataset, flag="data", loader_name="train_loader", + batch_size=BATCH, shuffle=True, is_training=True, + preload_labels=True) + crit = wl.watch_or_edit(nn.CrossEntropyLoss(reduction="none"), + flag="loss", signal_name=LOSS, per_sample=True, log=True) + + wl.serve(serving_grpc=False, serving_cli=False) + wl.start_training(timeout=0) + + tracker = wl.projection.get_tracker() + if tracker is None: + print("projection is DISABLED (WEIGHTSLAB_PROJECTION=0) -- no coordinates " + "will be written; the board stays hidden in the studio.") + else: + print(f"projection attached: every {tracker.every_n_steps} steps -> " + f"{tracker.out_dim}-D, columns {tracker.stats()['columns']}") + + print(f"\ntraining {CLUSTERS} clusters x {PER_CLUSTER} samples on {device}...") + for epoch in range(6): + for x, ids, y in loader: + x, y = x.to(device), y.to(device) + with wl.guard_training_context: + logits = model(x) + loss = crit(logits, y, batch_ids=ids, preds=logits.argmax(1, keepdim=True)) + opt.zero_grad() + loss.mean().backward() + opt.step() + if tracker is not None: + print(f" epoch {epoch + 1}: {tracker.steps_trained} fits, " + f"{tracker.samples_written} coordinates written, " + f"umap loss {tracker.last_loss:.4f}") + + if tracker is None: + return + + # ---- read the projection back exactly as the UI does ------------------- + # save_signals buffers; a live server sees the rows via the manager's own + # background flush thread, but this process is about to exit, so drain it. + from weightslab.backend.ledgers import get_dataframe + manager = get_dataframe() + manager.flush() + frame = manager._df + # Colour by the ground-truth class: the response then carries a per-point + # label aligned with the coordinates, so no id lookup is needed -- and it + # exercises the same colour path the board's "colour by" picker uses. + response = ps.build_projection_response( + frame, pb2.ProjectionRequest(max_points=100_000, color_column="target")) + print(f"\nGetProjection: success={response.success} :: {response.message}") + if not response.success: + return + + coords = np.array(response.coords, dtype=np.float32).reshape(-1, 3) + # Ground truth arrives aligned with the points (color_column="target"), so + # there is nothing to join: WeightsLab assigns its own sample uids and they + # are what the projection is keyed by, not the dataset's own index. + # Class ids come back as categorical LABELS ("0.0".."5.0"), not as a numeric + # ramp: a class id is a name, and ramping it would paint neighbouring + # classes near-identical shades. + labels = np.array(response.color_labels).astype(float).astype(np.int64) + print(f" {response.returned} points, extent " + f"x[{response.extent_min_x:.2f}, {response.extent_max_x:.2f}] " + f"y[{response.extent_min_y:.2f}, {response.extent_max_y:.2f}] " + f"z[{response.extent_min_z:.2f}, {response.extent_max_z:.2f}]") + + # Did the 3-D layout keep the structure? Compare the mean distance between + # cluster centroids against the mean spread within a cluster. + centroids = np.stack([coords[labels == c].mean(axis=0) for c in range(CLUSTERS)]) + within = float(np.mean([ + np.linalg.norm(coords[labels == c] - centroids[c], axis=1).mean() + for c in range(CLUSTERS)])) + between = float(np.mean([ + np.linalg.norm(centroids[i] - centroids[j]) + for i in range(CLUSTERS) for j in range(i + 1, CLUSTERS)])) + print(f" cluster separation: between/within = {between / max(within, 1e-9):.1f}x " + f"(within {within:.3f}, between {between:.3f})") + + # ---- re-projecting offline, from a different layer ---------------------- + # The live projection watched `classifier`'s input (the 16-D features). Say + # we decide that was the wrong layer: re-project from the trunk instead, + # with no training step, under its own prefix so both survive and the board + # can switch between them. + print("\nre-projecting offline from 'trunk.2' (no training step)...") + offline = wl.project_dataset(model, loader, layer="trunk.2", + prefix="umap_trunk", epochs=8, verbose=False) + print(f" {offline['samples']} samples x {offline['feature_dim']} features " + f"-> {offline['columns']} (loss {offline['final_loss']:.4f})") + + manager.flush() + frame = manager._df + print(f" projections now stored: {ps.available_prefixes(frame)}") + + alt = ps.build_projection_response( + frame, pb2.ProjectionRequest(prefix="umap_trunk", color_column="target")) + alt_xyz = np.array(alt.coords, dtype=np.float32).reshape(-1, 3) + alt_labels = np.array(alt.color_labels).astype(float).astype(np.int64) + alt_cent = np.stack([alt_xyz[alt_labels == c].mean(axis=0) for c in range(CLUSTERS)]) + alt_within = float(np.mean([ + np.linalg.norm(alt_xyz[alt_labels == c] - alt_cent[c], axis=1).mean() + for c in range(CLUSTERS)])) + alt_between = float(np.mean([ + np.linalg.norm(alt_cent[i] - alt_cent[j]) + for i in range(CLUSTERS) for j in range(i + 1, CLUSTERS)])) + print(f" 'umap_trunk' separation: {alt_between / max(alt_within, 1e-9):.1f}x " + f"(vs {between / max(within, 1e-9):.1f}x for the live 'umap')") + + # ---- bring your own projection ----------------------------------------- + # Any algorithm plugs into the same board: it draws every _x/_y(/_z) + # it finds. A 3-component PCA here (torch only, so the demo needs nothing + # extra); sklearn's TSNE(n_components=3) or umap.UMAP(n_components=3) drop + # in the same way, since anything with fit_transform is accepted. + def pca3(features): + centered = torch.from_numpy(features) + centered = centered - centered.mean(dim=0) + _, _, v = torch.pca_lowrank(centered, q=3) + return (centered @ v[:, :3]).numpy() + + print("\nyour own projection, on features WeightsLab collects (method=)...") + mine = wl.project_dataset(model, loader, method=pca3, prefix="pca", verbose=False) + print(f" {mine['samples']} samples x {mine['feature_dim']} features -> {mine['columns']}") + + # ...or compute everything yourself and hand over the coordinates: a PCA of + # the raw INPUTS, to set against what the model learned. + xs, uids = [], [] + for x, ids, _ in loader: + xs.append(x) + uids.extend(ids.tolist() if hasattr(ids, "tolist") else list(ids)) + raw = pca3(torch.cat(xs).numpy()) + print(f" wl.save_projection_coords -> " + f"{wl.save_projection_coords(raw, batch_ids=uids, prefix='pca_inputs')}") + + manager.flush() + frame = manager._df + print(f" projections now stored: {ps.available_prefixes(frame)}") + + # ---- the level-of-detail contract the 3-D board relies on -------------- + box = pb2.ProjectionRequest( + has_bounds=True, + min_x=float(coords[:, 0].min()), min_y=float(coords[:, 1].min()), + min_z=float(coords[:, 2].min()), + max_x=float(np.median(coords[:, 0])), max_y=float(coords[:, 1].max()), + max_z=float(coords[:, 2].max()), + max_points=200) + zoomed = ps.build_projection_response(frame, box) + print(f" zoomed box: {zoomed.returned} of {zoomed.total_in_view} in view " + f"(budget 200, {zoomed.total_available} projected overall)") + + +if __name__ == "__main__": + main() diff --git a/weightslab/integrations/ultralytics/trainer.py b/weightslab/integrations/ultralytics/trainer.py index 868cca16..bd0d05c6 100644 --- a/weightslab/integrations/ultralytics/trainer.py +++ b/weightslab/integrations/ultralytics/trainer.py @@ -42,6 +42,7 @@ import weightslab as wl from weightslab.backend import ledgers +from weightslab.projection import mirror_projection from .collate import wl_ul_dict_collate from .dataset import WLAwareDataset, WLAwareSegmentationDataset @@ -84,6 +85,22 @@ } +def _embedding_layer(model): + """The layer the live projection reads on a UL model: the last block before + the head (``model.model[-2]``, the P5 neck output for YOLO detect/segment), + whose output is pooled over H x W. Same layer UL's own ``YOLO.embed()`` + pools. ``None`` when *model* is not a UL layer stack (auto-pick then). + + Passed explicitly because the auto-pick ("input of the last conv") lands + in the head, on one scale's class branch at best -- and before frozen + layers were skipped, on the DFL box decoder, which training never runs. + """ + layers = getattr(model, "model", None) + if isinstance(layers, torch.nn.Sequential) and len(layers) >= 2: + return layers[-2] + return None + + class _WLTrainerMixin: """Shared WL callback wiring + dataloader construction for UL trainers. @@ -117,10 +134,17 @@ def _register_channels(): def _on_train_start(trainer): underlying = trainer.model trainer.optimizer = wl.watch_or_edit(trainer.optimizer, flag="optimizer") + embedding = _embedding_layer(underlying) trainer.model = wl.watch_or_edit( trainer.model, flag="model", forced_model_wrapping=True, compute_dependencies=False, + **({"projection": {"layer": embedding}} if embedding is not None else {}), ) + # Validation runs on UL's EMA shadow, which `_setup_train` built + # BEFORE this callback hooked the model -- so it carries no + # projection hook, and no val sample would ever be placed. + if getattr(trainer, "ema", None) is not None: + mirror_projection(trainer.ema.ema) # Get signals configuration from ledger signals_cfg = ledgers.get_hyperparams().get('signals_cfg', {}) diff --git a/weightslab/opencode_process.py b/weightslab/opencode_process.py index 92db1eee..135f2834 100644 --- a/weightslab/opencode_process.py +++ b/weightslab/opencode_process.py @@ -223,6 +223,25 @@ def lock_path(workspace_dir: str) -> Path: return Path(workspace_dir) / LOCK_FILENAME +def machine_lock_path() -> Path: + """Machine-wide fallback lock, for sides that do NOT share a workspace. + + The per-workspace lock is the rendezvous when both halves agree on a + directory — `weightslab start ` and a backend whose + WEIGHTSLAB_ROOT_LOG_DIR is that same dir. They frequently do not: the UI + server generates its own workspace when started without one, while the + training backend falls back to its cwd. Each then found no lock, spawned + its OWN OpenCode, and the two sides ran different models with no error + anywhere — the studio's model picker wrote one server's config while the + backend read the other's. + + Machine-wide rather than per-user-per-project on purpose: OpenCode is a + single local server by design, and OPENCODE_URL still overrides for anyone + who genuinely wants several. + """ + return Path.home() / ".weightslab" / LOCK_FILENAME + + def read_lock(workspace_dir: str) -> Optional[dict]: try: with open(lock_path(workspace_dir), "r", encoding="utf-8") as f: @@ -234,11 +253,33 @@ def read_lock(workspace_dir: str) -> Optional[dict]: return data +def read_machine_lock() -> Optional[dict]: + """The machine-wide lock, or None when absent/unreadable.""" + try: + with open(machine_lock_path(), "r", encoding="utf-8") as f: + data = json.load(f) + except (OSError, ValueError): + return None + if not isinstance(data, dict) or not data.get("url"): + return None + return data + + def write_lock(workspace_dir: str, url: str, pid: Optional[int] = None) -> None: + payload = {"url": url, "pid": pid} + # Publish machine-wide too, so a side that does not share this workspace + # can still find this server instead of spawning a second one. + try: + machine = machine_lock_path() + os.makedirs(machine.parent, exist_ok=True) + with open(machine, "w", encoding="utf-8") as f: + json.dump(payload, f) + except OSError: + _LOGGER.debug("Could not write the machine-wide OpenCode lock") try: os.makedirs(workspace_dir, exist_ok=True) with open(lock_path(workspace_dir), "w", encoding="utf-8") as f: - json.dump({"url": url, "pid": pid}, f) + json.dump(payload, f) except OSError: # Best-effort: a later caller simply won't find this server via the # lock file and will spawn its own instead. Not worth failing the @@ -432,6 +473,20 @@ def resolve_or_spawn_opencode(workspace_dir: str, origin: Optional[str] = None, ) return {"ok": True, "url": lock["url"], "source": "lockfile"} + # Nothing for THIS workspace: before spawning a second server, adopt one + # already running on this machine. Without this the UI and the training + # backend — which routinely disagree about the workspace — each ran their + # own, so a model chosen in one was invisible to the other. + shared = read_machine_lock() + if shared and opencode_healthy(shared["url"]): + _LOGGER.info( + "OpenCode: adopting the server at %s (machine-wide lock) for " + "workspace %s.", shared["url"], workspace_dir, + ) + # Re-publish under this workspace so the next lookup here is direct. + write_lock(workspace_dir, shared["url"], shared.get("pid")) + return {"ok": True, "url": shared["url"], "source": "machine-lockfile"} + argv = resolve_opencode_argv() if argv is None: return { @@ -454,6 +509,10 @@ def resolve_or_spawn_opencode(workspace_dir: str, origin: Optional[str] = None, try: process = subprocess.Popen( cmd, cwd=workspace_dir, + # stdin too: a server never reads it, and inheriting the parent's + # fails on Windows when that has none usable (a Jupyter kernel, + # pythonw, a detached task) -- WinError 6 / 50 at spawn. + stdin=subprocess.DEVNULL, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL, # Detached: this server is meant to outlive whichever side # happened to spawn it (the other side may still be using it diff --git a/weightslab/projection/__init__.py b/weightslab/projection/__init__.py new file mode 100644 index 00000000..d6f2aea3 --- /dev/null +++ b/weightslab/projection/__init__.py @@ -0,0 +1,76 @@ +"""Live low-dimensional projections of a model's representation. + +Built in: parametric UMAP (:mod:`parametric_umap`), which trains a small +``features -> R^3`` encoder alongside the model and writes each sample's +coordinates back as ordinary per-sample signals. + +Your own: any projection you compute (t-SNE, umap-learn, PCA...) plugs into the +same board through :func:`save_projection_coords`, or through +``project_dataset(..., method=...)`` (see :mod:`custom`). +""" + +from weightslab.projection.parametric_umap import ( # noqa: F401 + ENV_ENABLED, + ENV_DIM, + ENV_EVERY, + ENV_GRAPH, + ENV_NEIGHBORS, + ProjectionEncoder, + ProjectionTracker, + attach_projection, + clear_pending_restore, + detach_projection, + encoder_state, + find_ab_params, + first_tensor, + flatten_features, + checkpoint_path, + get_tracker, + layer_detail, + load_projection, + membership_high_dim, + mirror_projection, + observe_batch, + pick_embedding_layer, + project_dataset, + projection_enabled, + reattach_projection, + resolve_module, + restore_encoder, + restore_from_checkpoint, + save_projection, +) +from weightslab.projection.custom import save_projection_coords # noqa: F401 + +__all__ = [ + "ENV_ENABLED", + "ENV_DIM", + "ENV_EVERY", + "ENV_GRAPH", + "ENV_NEIGHBORS", + "ProjectionEncoder", + "ProjectionTracker", + "attach_projection", + "clear_pending_restore", + "detach_projection", + "encoder_state", + "find_ab_params", + "first_tensor", + "flatten_features", + "checkpoint_path", + "get_tracker", + "layer_detail", + "load_projection", + "membership_high_dim", + "mirror_projection", + "observe_batch", + "pick_embedding_layer", + "project_dataset", + "projection_enabled", + "reattach_projection", + "resolve_module", + "restore_encoder", + "restore_from_checkpoint", + "save_projection", + "save_projection_coords", +] diff --git a/weightslab/projection/custom.py b/weightslab/projection/custom.py new file mode 100644 index 00000000..979d4575 --- /dev/null +++ b/weightslab/projection/custom.py @@ -0,0 +1,189 @@ +"""Projections you compute yourself: t-SNE, umap-learn, PCA, anything. + +The Projection Board draws any per-sample ``signals//_x`` / +``_y`` (/ ``_z``) columns, whatever wrote them, so plugging your own algorithm +in is a matter of writing coordinates under a prefix of your choosing: + +* :func:`save_projection_coords` -- you have the coordinates (and the sample + ids they belong to); WeightsLab stores them. +* ``wl.project_dataset(model, loader, method=...)`` -- WeightsLab collects the + features from the model for you and hands them to your ``method``. + +Either way the result sits beside the built-in parametric UMAP (``umap``), and +the board's picker switches between them. +""" + +from __future__ import annotations + +import logging +import re + +import numpy as np + +logger = logging.getLogger(__name__) + +# Becomes part of a column name (`signals//_x`), so keep it to what +# survives every storage layer and query: no "/" (the signals// separator), no +# spaces, nothing that needs quoting in a query. +_PREFIX_RE = re.compile(r"^[A-Za-z0-9][A-Za-z0-9_.-]*$") +AXES = ("x", "y", "z") + + +def validate_prefix(prefix) -> str: + """*prefix*, checked. Raises ``ValueError`` with what is allowed.""" + name = str(prefix or "").strip() + if not _PREFIX_RE.match(name): + raise ValueError( + f"projection prefix {prefix!r} is not usable as a column name: use " + f"letters, digits, '_', '.' or '-', starting with a letter or digit " + f"(e.g. 'tsne', 'umap_learn', 'pca.layer3').") + return name + + +def as_coordinates(coords, n: int | None = None) -> np.ndarray: + """Coerce *coords* to a float32 ``(N, 2)`` or ``(N, 3)`` array. + + Accepts a torch tensor, a numpy array, or a nested list. Rows that are + NaN/inf are allowed -- they mean "not placed" and the board skips them. + """ + if hasattr(coords, "detach"): + coords = coords.detach().cpu().numpy() + array = np.asarray(coords, dtype=np.float32) + if array.ndim != 2 or array.shape[1] not in (2, 3): + raise ValueError( + f"projection coordinates must be shaped (N, 2) or (N, 3); got " + f"{tuple(array.shape)}. Reduce to 2 or 3 components first " + f"(e.g. TSNE(n_components=3)).") + if n is not None and array.shape[0] != n: + raise ValueError( + f"projection coordinates have {array.shape[0]} rows but there are " + f"{n} sample ids; they must be aligned one to one.") + return array + + +def method_name(method) -> str: + """A prefix for *method* when the caller did not name one. + + ``TSNE(...)`` -> ``"tsne"``, ``def my_pca(x)`` -> ``"my_pca"``; a lambda, + having no useful name, -> ``"custom"``. + """ + name = getattr(method, "__name__", None) or type(method).__name__ + if name.startswith("<"): + return "custom" + name = re.sub(r"[^A-Za-z0-9_.-]", "_", name.lower()).strip("_.-") + return name or "custom" + + +def run_method(method, features: np.ndarray) -> np.ndarray: + """Run a user projection on ``(N, F)`` features, return ``(N, 2|3)``. + + *method* is anything with a scikit-learn style ``fit_transform`` (``TSNE``, + ``PCA``, ``umap.UMAP``) or a plain callable ``features -> coords``. + """ + fit_transform = getattr(method, "fit_transform", None) + if callable(fit_transform): + out = fit_transform(features) + elif callable(method): + out = method(features) + else: + raise TypeError( + f"method must be 'umap', a callable features -> coords, or an object " + f"with fit_transform (e.g. sklearn's TSNE); got {type(method).__name__}.") + return as_coordinates(out, n=features.shape[0]) + + +def check_not_live_prefix(prefix: str) -> None: + """Refuse to write under the prefix the live parametric UMAP owns. + + The live projection rewrites its prefix on every fit, a few hundred samples + at a time, so a user projection stored there would be overwritten + piecemeal: one cloud holding two unrelated layouts. Turn the built-in off + (``projection=False``) to reuse its name. + """ + from weightslab.projection.parametric_umap import get_tracker + live = get_tracker() + if live is not None and live._handle is not None and live.signal_prefix == prefix: + raise ValueError( + f"prefix {prefix!r} belongs to the live parametric UMAP, which keeps " + f"rewriting it during training. Pick another prefix (e.g. 'tsne'), or " + f"turn the built-in off with wl.watch_or_edit(model, flag='model', " + f"projection=False).") + + +def save_projection_coords(coords, batch_ids, prefix: str) -> list: + """Store a projection you computed yourself, for the Projection Board. + + Writes ``signals//_x``, ``_y`` (and ``_z`` for 3-D) onto each + sample's row -- ordinary per-sample signals, so the grid can sort, filter + and export them too -- and the board lists *prefix* in its picker beside + the built-in ``umap``. + + The samples are NOT counted as seen: placing a sample in a picture is not + the model training on it, so ``nb_seen`` / ``last_seen`` are left alone. + + Call it whenever you like -- at step 500, once per epoch, after training. + Writing the same prefix again replaces the coordinates of the samples you + pass; samples you leave out keep theirs. + + Args: + coords: ``(N, 2)`` or ``(N, 3)`` -- tensor, array or nested list. A 2-D + projection is drawn on the z = 0 plane. NaN rows are skipped. + batch_ids: the ``N`` sample ids the rows belong to, in the same order + (the ids the WeightsLab loaders yield). + prefix: the projection's name, e.g. ``"tsne"``. Not ``"umap"`` while the + built-in live projection is running (it would overwrite yours). + + Returns: + list: the column names written. + + Example: + >>> from sklearn.manifold import TSNE + >>> coords = TSNE(n_components=3).fit_transform(features) # (N, 3) + >>> wl.save_projection_coords(coords, batch_ids=ids, prefix="tsne") + """ + from weightslab import src as _src + + prefix = validate_prefix(prefix) + check_not_live_prefix(prefix) + if hasattr(batch_ids, "detach"): + ids = [str(i) for i in batch_ids.detach().cpu().tolist()] + else: + ids = [str(i) for i in batch_ids] + array = as_coordinates(coords, n=len(ids)) + if not ids: + return [] + + signals = {f"{prefix}_{axis}": array[:, i] for i, axis in enumerate(AXES[:array.shape[1]])} + if array.shape[1] == 2: + # A prefix that once held a 3-D projection still has its _z column, and + # the board reads x, y AND z when all three exist -- new x/y would be + # drawn against stale depth. Flatten it rather than leave it. + try: + # The handle save_signals writes through, so both agree on which + # dataframe "already has a _z column" refers to. + frame_m = _src.DATAFRAME_M if _src.DATAFRAME_M is not None else _src.get_dataframe() + if f"signals//{prefix}_z" in frame_m.get_df_view(limit=1).columns: + signals[f"{prefix}_z"] = np.zeros(len(ids), dtype=np.float32) + except Exception: + pass + + _src.save_signals(signals=signals, batch_ids=ids, log=False, _seen=False) + from weightslab.projection.registry import register_prefix + register_prefix(prefix) + # Flush now. save_signals only buffers, and the board reads flushed rows: + # with the usual 30 s flush interval, the cell that just "plugged in" a + # t-SNE was followed by half a minute of "No projection found". The live + # projection writes every step and rightly stays buffered; this is one + # explicit write the user is about to go and look at. + try: + frame_m = _src.DATAFRAME_M if _src.DATAFRAME_M is not None else _src.get_dataframe() + if hasattr(frame_m, "flush"): + frame_m.flush() + except Exception as exc: + logger.debug(f"[projection] '{prefix}': immediate flush failed ({exc}); " + f"the board shows it at the next scheduled flush") + written = [f"signals//{name}" for name in signals] + placed = int(np.isfinite(array).all(axis=1).sum()) + logger.info(f"[projection] '{prefix}': {placed} of {len(ids)} samples placed " + f"({array.shape[1]}-D) -> {written}") + return written diff --git a/weightslab/projection/parametric_umap.py b/weightslab/projection/parametric_umap.py new file mode 100644 index 00000000..e8ff3410 --- /dev/null +++ b/weightslab/projection/parametric_umap.py @@ -0,0 +1,1773 @@ +"""Parametric UMAP: a live 3-D projection of what the model is learning. + +A small encoder network ``features -> R^3`` is trained *alongside* your model, +on the UMAP cross-entropy, from the activations of one layer. Unlike +``umap-learn``'s ``fit_transform`` this is **parametric**: the projection is a +function, so a sample's coordinates come from one forward pass and the +embedding stays comparable from step to step instead of being re-laid-out from +scratch on every refit. + +Why parametric, here: + * it runs online during training, on batches the model already computed, so + there is no separate embedding pass over the dataset; + * new/unseen samples project without refitting; + * the embedding moves smoothly as the representation moves, which is the + thing worth watching -- a fresh UMAP per epoch mostly shows you UMAP's own + initialisation noise. + +It NEVER perturbs training: features are detached at the hook and the encoder +carries its own optimizer. Turning it off (see :func:`projection_enabled`) +removes the hook entirely. + +Algorithm (standard UMAP objective, computed densely per batch): + 1. kNN graph on the batch's features, with UMAP's per-point ``rho`` + (distance to nearest neighbour) and ``sigma`` (binary-searched so the + row's memberships sum to ``log2(k)``). + 2. High-dim membership ``mu_ij = exp(-(d_ij - rho_i) / sigma_i)``, + symmetrised as ``mu + mu^T - mu*mu^T``. + 3. Low-dim membership ``nu_ij = 1 / (1 + a * d_ij^(2b))``, with ``(a, b)`` + fitted from ``spread``/``min_dist`` as umap-learn does. + 4. Cross-entropy between the two: attractive ``-mu*log(nu)`` on the kNN + edges, repulsive ``-(1-mu)*log(1-nu)`` on the rest. + +Batch-level density (B up to a few hundred) keeps this a handful of B x B +matmuls -- cheap next to the model's own step, and it only runs every +``every_n_steps`` steps. +""" + +from __future__ import annotations + +import contextlib +import logging +import os +import re +import threading + +import numpy as np +import torch as th +import torch.nn as nn +from tqdm import tqdm + +logger = logging.getLogger(__name__) + +# Env switches. The feature is ON by default; set WEIGHTSLAB_PROJECTION to one +# of 0/false/no/off to remove it completely (no hook, no encoder, no cost). +ENV_ENABLED = "WEIGHTSLAB_PROJECTION" +ENV_EVERY = "WEIGHTSLAB_PROJECTION_EVERY" +ENV_DIM = "WEIGHTSLAB_PROJECTION_DIM" +ENV_NEIGHBORS = "WEIGHTSLAB_PROJECTION_NEIGHBORS" +ENV_GRAPH = "WEIGHTSLAB_PROJECTION_GRAPH" + +_OFF_VALUES = {"0", "false", "no", "off"} + +# How many fit-visits with no captured features before we say so out loud. +_MISS_WARN_AFTER = 3 +# Live encoder is checkpointed this often (in fits). +_SAVE_EVERY_FITS = 20 +# Consecutive failed fits before the projection gives up for the run. One +# failure is logged and retried at the next fit; a streak this long is a cause +# that will not go away by itself (an OOM, a feature shape the encoder cannot +# take), and every retry costs a kNN graph inside the user's training step. +_MAX_CONSECUTIVE_FAILURES = 5 + + +def projection_enabled() -> bool: + """Whether live projection is on. Default ON; ``WEIGHTSLAB_PROJECTION=0`` + (or false/no/off, any case) turns it off.""" + raw = os.environ.get(ENV_ENABLED) + if raw is None or str(raw).strip() == "": + return True + return str(raw).strip().lower() not in _OFF_VALUES + + +def _env_int(name: str, default: int) -> int: + try: + value = int(str(os.environ.get(name, "")).strip()) + return value if value > 0 else default + except (TypeError, ValueError): + return default + + +def find_ab_params(spread: float = 1.0, min_dist: float = 0.1, + iters: int = 400) -> tuple[float, float]: + """Fit ``(a, b)`` so ``1/(1 + a*d^(2b))`` approximates UMAP's piecewise + target curve -- the same fit umap-learn does with ``scipy.curve_fit``, done + here with a few hundred Adam steps so this module needs no scipy. + + Falls back to the well-known ``spread=1.0, min_dist=0.1`` solution if the + fit misbehaves, so a bad input can never leave the caller without params. + """ + try: + xv = th.linspace(0.0, spread * 3.0, 300) + yv = th.where(xv < min_dist, th.ones_like(xv), + th.exp(-(xv - min_dist) / spread)) + log_a = th.zeros(1, requires_grad=True) + log_b = th.zeros(1, requires_grad=True) + opt = th.optim.Adam([log_a, log_b], lr=0.05) + for _ in range(iters): + opt.zero_grad() + pred = 1.0 / (1.0 + log_a.exp() * xv.pow(2.0 * log_b.exp())) + loss = ((pred - yv) ** 2).mean() + loss.backward() + opt.step() + a = float(log_a.exp().item()) + b = float(log_b.exp().item()) + if not (np.isfinite(a) and np.isfinite(b)) or a <= 0 or b <= 0: + raise ValueError("degenerate fit") + return a, b + except Exception: + # umap-learn's values for spread=1.0 / min_dist=0.1. + return 1.5769434, 0.8950608 + + +class ProjectionEncoder(nn.Module): + """The parametric half of parametric UMAP: ``features -> R^out_dim``. + + Deliberately small (two hidden layers). It has to keep up with a moving + representation on a per-batch budget, and a bigger encoder mostly buys + overfitting to whatever the last few batches looked like. + """ + + def __init__(self, in_dim: int, out_dim: int = 3, hidden: int = 256): + super().__init__() + self.net = nn.Sequential( + nn.Linear(in_dim, hidden), + nn.ReLU(inplace=True), + nn.Linear(hidden, hidden // 2), + nn.ReLU(inplace=True), + nn.Linear(hidden // 2, out_dim), + ) + + def forward(self, x: th.Tensor) -> th.Tensor: + return self.net(x) + + +def membership_high_dim(feats: th.Tensor, n_neighbors: int, + bisect_steps: int = 32) -> th.Tensor: + """UMAP's fuzzy simplicial set for one batch, as a dense ``(B, B)`` matrix. + + ``rho_i`` is i's distance to its nearest neighbour and ``sigma_i`` is + binary-searched so that ``sum_j exp(-(d_ij - rho_i)/sigma_i) == log2(k)``. + That per-point normalisation is the whole reason UMAP tolerates wildly + varying local density -- without it, dense regions dominate the loss. + """ + n = feats.shape[0] + k = max(2, min(int(n_neighbors), n - 1)) + + dist = th.cdist(feats, feats) + dist = dist + th.eye(n, device=feats.device) * 1e12 # exclude self + + knn_d, knn_i = th.topk(dist, k, dim=1, largest=False) + rho = knn_d[:, 0:1] # (B,1) + centered = (knn_d - rho).clamp(min=0.0) + + target = float(np.log2(k)) + lo = th.full((n, 1), 1e-6, device=feats.device) + hi = th.full((n, 1), 1e4, device=feats.device) + for _ in range(bisect_steps): + mid = (lo + hi) / 2.0 + total = th.exp(-centered / mid).sum(dim=1, keepdim=True) + too_big = total > target + hi = th.where(too_big, mid, hi) + lo = th.where(too_big, lo, mid) + sigma = ((lo + hi) / 2.0).clamp(min=1e-6) + + mu = th.zeros(n, n, device=feats.device) + mu.scatter_(1, knn_i, th.exp(-centered / sigma)) + # Fuzzy union: an edge counts if EITHER endpoint considers it a neighbour. + return mu + mu.t() - mu * mu.t() + + +def first_tensor(payload): + """The first tensor inside whatever a module handed us. + + Real models return far more than a bare tensor: tuples + (``(logits, hidden)``), HuggingFace ``ModelOutput`` dicts, detection heads + returning lists of per-level maps. Reaching the tensor is the difference + between "works on a toy CNN" and "works on the model the user actually has", + so unwrap one level of the usual containers rather than giving up. + """ + if isinstance(payload, th.Tensor): + return payload + if isinstance(payload, dict): + for key in ("last_hidden_state", "hidden_states", "features", "logits"): + value = payload.get(key) + if isinstance(value, th.Tensor): + return value + payload = list(payload.values()) + if isinstance(payload, (tuple, list)): + for item in payload: + found = first_tensor(item) + if found is not None: + return found + return None + + +def flatten_features(tensor: th.Tensor, feature_last: bool) -> th.Tensor: + """Reduce an activation of any rank to ``(B, F)``. + + WHICH axis holds the features depends on what was hooked, and getting this + backwards silently projects the wrong thing: + + * ``feature_last=True`` -- the INPUT of an ``nn.Linear``. Linear always acts + on the last axis, so that axis is the feature axis whatever the rank: + ``(B, T, C)`` from a transformer pools over the tokens ``T`` and keeps the + model dim ``C``. Pooling the other way round would hand the projection a + per-token scalar and call it an embedding. + * ``feature_last=False`` -- the OUTPUT of a conv/norm/pool, which is + channels-first: ``(B, C, H, W)`` pools over the spatial axes and keeps + ``C``. + + Both collapse to ``(B, F)``; a tensor already 2-D is returned unchanged. + """ + if tensor.ndim < 2: + return tensor + if tensor.ndim == 2: + return tensor + if feature_last: + return tensor.mean(dim=tuple(range(1, tensor.ndim - 1))) + return tensor.mean(dim=tuple(range(2, tensor.ndim))) + + +def _inert_hook(_module, _inputs, _output): + """What a projection hook becomes when its model is serialised: nothing. + + Module-level so pickle stores it by reference -- a few bytes, and no state. + """ + return None + + +class _FeatureHook: + """The forward hook, as an object rather than a closure. + + A closure over the tracker is what the hook used to be, and dill (which + saves WeightsLab's model architectures) serialises closures BY VALUE: every + architecture file carried a full copy of the tracker -- encoder, optimizer, + the feature buffer -- and a model restored from one came back with a hook + feeding that dead copy, so the live projection never saw its features + again. Pickled, this object is :func:`_inert_hook` instead. + + ``copy.deepcopy`` keeps the SAME hook, so a deep copy of the model (an EMA + shadow, say) keeps feeding the live tracker exactly as the closure did. + """ + + __slots__ = ("tracker", "use_input", "feature_last") + + def __init__(self, tracker, use_input: bool, feature_last: bool): + self.tracker = tracker + self.use_input = use_input + self.feature_last = feature_last + + def __call__(self, module, inputs, output): + try: + tensor = first_tensor(inputs if self.use_input else output) + if tensor is None or tensor.ndim < 2: + return + self.tracker._pending = flatten_features( + tensor.detach().float(), feature_last=self.feature_last) + self.tracker._pending_training = bool(getattr(module, "training", True)) + except Exception as exc: # never break a forward pass + logger.debug(f"[projection] feature hook failed: {exc}") + + def __reduce__(self): + return (_return_inert_hook, ()) + + def __deepcopy__(self, memo): + return self + + +def _return_inert_hook(): + return _inert_hook + + +def _is_stale_hook(hook) -> bool: + """A projection hook that no live tracker owns. + + The inert placeholder, and the closure earlier versions installed: an + architecture pickled with one of those restores with a hook feeding a dead + copy of the tracker, and would go on capturing every forward for nothing. + """ + if hook is _inert_hook: + return True + return getattr(hook, "__qualname__", "").startswith("ProjectionTracker.attach.") + + +def _purge_inert_hooks(model: nn.Module) -> None: + """Drop the stale projection hooks a deserialised model carries. + + The inert ones are harmless where they are, but each save/restore cycle + would add one more (see _FeatureHook). + """ + try: + for module in model.modules(): + hooks = getattr(module, "_forward_hooks", None) + if not hooks: + continue + for key in [k for k, h in hooks.items() if _is_stale_hook(h)]: + del hooks[key] + except Exception as exc: + logger.debug(f"[projection] could not purge inert hooks: {exc}") + + +class ProjectionTracker: + """Owns the encoder, the feature hook and the write-back of coordinates. + + Lifecycle, per tracked step:: + + forward hook (detached features) + -> observe_batch(ids, step) # from the loss wrapper, where the + # sample ids for that same forward + # pass are known + -> fit the encoder on the batch + -> write umap_x / umap_y / umap_z per sample + + Pairing features with ids positionally at the loss is the same trick + ``ctx.logits`` already uses: the hook cannot see sample ids, the loss cannot + see intermediate activations, and both describe the same batch in the same + order. + """ + + def __init__(self, out_dim: int = 3, n_neighbors: int = 15, + every_n_steps: int = 50, inner_steps: int = 4, + lr: float = 1e-3, min_dist: float = 0.1, spread: float = 1.0, + repulsion_strength: float = 1.0, signal_prefix: str = "umap", + graph_size: int | None = None): + self.out_dim = int(out_dim) + self.n_neighbors = int(n_neighbors) + self.every_n_steps = max(1, int(every_n_steps)) + self.inner_steps = max(1, int(inner_steps)) + self.lr = float(lr) + self.repulsion_strength = float(repulsion_strength) + self.signal_prefix = str(signal_prefix) + + self.a, self.b = find_ab_params(spread=spread, min_dist=min_dist) + + # Rolling buffer of recent (features, ids), fitted on instead of the + # bare batch. + # + # This is not an optimisation, it is a correctness fix. UMAP's graph is + # a k-NEAREST-neighbour graph, so with a batch of 16 and k=15 every + # sample is every other sample's neighbour: the graph is complete, the + # attractive term pulls on all pairs equally and the projection has no + # local structure to preserve. Measured on clustered data, batch 16 + # separates within- from between-cluster membership by 3x; batch 512 + # by ~10^11. Buffering decouples graph quality from whatever batch size + # the user's training loop happens to use. + self.graph_size = int(graph_size or _env_int(ENV_GRAPH, 512)) + self._buffer_feats: list = [] + self._buffer_ids: list = [] + + self._encoder: ProjectionEncoder | None = None + self._optimizer: th.optim.Optimizer | None = None + # How the hook reads the layer -- (use_input, feature_last) -- kept so + # a mirror on another copy of that layer reads it the same way. + self._hook_spec = (True, True) + # Hooks on copies of the layer in OTHER models (see mirror). + self._mirror_handles: list = [] + # Feature width the encoder was built for. Compared on every fit so a + # live architecture edit rebuilds it instead of raising forever. + self._encoder_dim: int | None = None + self._previous_dim: int | None = None + self.rebuilds = 0 + self._handle = None + self._pending: th.Tensor | None = None + # Whether the hooked module was in train mode when it produced + # _pending. Eval batches are PLACED by the current encoder but never + # FIT on: held-out samples belong in the picture, and letting them + # shape the projection would leak the eval set into it. + self._pending_training = True + self._lock = threading.Lock() + # save_signals is how coordinates get written, and save_signals also + # calls back into observe_batch -- without this the write-back would + # recurse into another fit. + self._writing_back = False + self._last_step_run = -1 + # Step of the last batch whose features were taken (or found missing), + # so a second per-sample write for the same batch is not a "miss". + self._last_step_seen = -1 + self._buffer_count = 0 + self._missed_batches = 0 + self.steps_trained = 0 + self.samples_written = 0 + # The fit count at the last encoder checkpoint (see _write_coords). + self._last_saved_fit = 0 + self.last_loss: float | None = None + # Which layer is hooked, as a named_modules() name of the model's root + # module, so a checkpoint restore that swaps in a NEW model object can + # hook the same layer on it (see reattach_projection). + self.layer_name: str | None = None + # What the hook reads there, e.g. "C3k2 output" (see layer_detail). + self.layer_detail: str | None = None + self._registered_prefix: str | None = None + # Failure accounting (see _record_failure). The projection is + # instrumentation: it may stop, it may never take the run down. + self.failures = 0 + self._consecutive_failures = 0 + self.disabled_reason: str | None = None + self.skipped_nonfinite = 0 + # Batches whose hooked features did not have one row per sample, and the + # last such (rows, samples) -- see observe_batch. + self.mismatched_batches = 0 + self.last_mismatch: tuple[int, int] | None = None + + # ------------------------------------------------------------------ hook + def attach(self, module: nn.Module, use_input: bool = True, + feature_last: bool | None = None): + """Hook *module*. + + ``use_input`` captures what flows INTO it (the representation) rather + than its output (usually logits). ``feature_last`` says which axis holds + the features -- ``True`` for a Linear's input, ``False`` for a conv's + channels-first tensor. It defaults to ``use_input`` only to keep the + two-argument call working; :func:`pick_embedding_layer` always states it. + """ + # Only the primary hook: mirrors sit on other models that are still + # running (an EMA shadow), and a re-attach to a restored model must + # not silence them. + self._remove_handle() + if feature_last is None: + feature_last = use_input + self._hook_spec = (bool(use_input), bool(feature_last)) + self._handle = module.register_forward_hook( + _FeatureHook(self, *self._hook_spec)) + return self._handle + + def mirror(self, module: nn.Module) -> bool: + """Also capture from *module*: the hooked layer's twin in another copy + of the model, which then feeds this tracker exactly like the original. + + For a copy made BEFORE the projection attached -- one made after + already carries the hook (see ``_FeatureHook.__deepcopy__``). Idempotent. + """ + if self._handle is None: + return False + hooks = getattr(module, "_forward_hooks", None) or {} + if any(isinstance(h, _FeatureHook) and h.tracker is self for h in hooks.values()): + return True + self._mirror_handles.append( + module.register_forward_hook(_FeatureHook(self, *self._hook_spec))) + return True + + def _remove_handle(self) -> None: + if self._handle is not None: + try: + self._handle.remove() + except Exception: + pass + self._handle = None + + def detach(self) -> None: + self._remove_handle() + for handle in self._mirror_handles: + try: + handle.remove() + except Exception: + pass + self._mirror_handles = [] + + def reset(self) -> None: + """Forget the learned layout: no encoder, empty buffer, no fits. + + What a checkpoint restore lands on when that checkpoint was taken + before the first fit -- the honest state of the projection at that + moment is "nothing learned yet", not whatever a later step learned. + """ + with self._lock: + self._encoder = None + self._optimizer = None + self._encoder_dim = None + self._previous_dim = None + self._pending = None + self._clear_buffer() + self.steps_trained = 0 + self._last_saved_fit = 0 + self.last_loss = None + + def _record_failure(self, exc: Exception, step: int) -> None: + """Count a failed fit; warn once; give up after a streak. + + Training never sees any of this. Before, every failure was a DEBUG + line, so a projection that failed on every fit looked exactly like one + that was simply slow to appear. + """ + self.failures += 1 + self._consecutive_failures += 1 + if self.failures == 1: + logger.warning( + f"[projection] fit failed at step {step}: {exc!r}. Training is " + f"unaffected; the projection retries at its next fit.") + else: + logger.debug(f"[projection] fit failed at step {step}: {exc!r}") + if self._consecutive_failures >= _MAX_CONSECUTIVE_FAILURES: + self.disabled_reason = ( + f"{self._consecutive_failures} consecutive failed fits " + f"(last at step {step}: {exc!r})") + self.detach() + logger.warning( + f"[projection] turned off for the rest of this run after " + f"{self.disabled_reason}. Training continues and the coordinates " + f"already written stay on the board. Set {ENV_ENABLED}=0, or " + f"projection=False on the model, to skip it from the start.") + + # ------------------------------------------------------------- training + def _ensure_encoder(self, in_dim: int, device) -> None: + """Build the encoder, and REBUILD it when the feature width changes. + + This is not a defensive nicety in WeightsLab: editing the architecture + of a running model -- pruning neurons, growing a layer -- is a headline + feature, and it changes the hooked layer's width underneath us. An + encoder pinned to the old width then raises on every batch forever, + which the caller's broad ``except`` turns into a projection that + silently stops updating after the user's first edit. + + Rebuilding loses the learned layout, so the cloud visibly re-lays-out. + That is the honest outcome -- the representation being projected really + did change -- and it is logged at INFO so the jump has an explanation. + """ + if self._encoder is not None and self._encoder_dim == in_dim: + # The model may also have been moved (.cuda() after wrapping), or + # the encoder may have been adopted from an offline fit that ran on + # another device. Follow it rather than dying on a mismatch. + if next(self._encoder.parameters()).device != device: + self._encoder.to(device) + # Adam's state (exp_avg, exp_avg_sq) does NOT follow .to(), so + # step() would then mix devices and raise on every batch -- + # which the caller's broad except turns into a projection that + # silently stops. Rebuild rather than migrate the state: a few + # steps of lost momentum is nothing next to that. + self._optimizer = th.optim.Adam(self._encoder.parameters(), lr=self.lr) + return + + rebuilt = self._encoder is not None + if rebuilt: + # Buffered features describe the OLD layer width; mixing them into + # the next graph would be comparing different spaces. + self._clear_buffer() + self._encoder = ProjectionEncoder(in_dim, self.out_dim).to(device) + self._optimizer = th.optim.Adam(self._encoder.parameters(), lr=self.lr) + self._encoder_dim = in_dim + self.rebuilds += 1 if rebuilt else 0 + if rebuilt: + logger.info( + f"[projection] feature width changed {self._previous_dim} -> {in_dim} " + f"(architecture edit?); encoder rebuilt, so the layout restarts" + ) + else: + logger.info( + f"[projection] parametric UMAP encoder {in_dim} -> {self.out_dim} " + f"(a={self.a:.4f}, b={self.b:.4f}, every {self.every_n_steps} steps)" + ) + self._previous_dim = in_dim + + def _umap_loss(self, feats: th.Tensor, target_p: th.Tensor) -> th.Tensor: + emb = self._encoder(feats) + d2 = th.cdist(emb, emb).pow(2).clamp(min=1e-12) + nu = 1.0 / (1.0 + self.a * d2.pow(self.b)) + nu = nu.clamp(1e-6, 1.0 - 1e-6) + + off_diag = 1.0 - th.eye(feats.shape[0], device=feats.device) + attractive = -(target_p * th.log(nu)) * off_diag + repulsive = -((1.0 - target_p) * th.log(1.0 - nu)) * off_diag + return (attractive.sum() + + self.repulsion_strength * repulsive.sum()) / off_diag.sum() + + def _push_buffer(self, feats: th.Tensor, ids: list) -> None: + """Add one training batch to the rolling buffer. + + Called for EVERY training batch, not only on fit steps. It used to be + fed only the batch that happened to land on a fit step -- one batch in + ``every_n_steps`` -- so the "most recent graph_size samples" were really + the last few fit steps' batches, spread over hundreds of steps of + representation drift, and samples that never landed on a fit step + never got coordinates at all (about a quarter of them in the bundled + demo). Recency matters: the representation is moving, so features from + long ago describe a model that no longer exists. + + Cheap per step: a list append. Whole chunks are dropped from the front + once the rest still fills the graph, so memory stays at graph_size plus + one batch; the single concatenation happens at fit time. + """ + if self._buffer_feats and self._buffer_feats[-1].shape[1] != feats.shape[1]: + # The hooked layer changed width (an architecture edit): buffered + # features live in the old space, and mixing them would compare + # different spaces. _ensure_encoder rebuilds at the next fit. + self._clear_buffer() + self._buffer_feats.append(feats) + self._buffer_ids.extend(ids) + self._buffer_count += int(feats.shape[0]) + while (len(self._buffer_feats) > 1 + and self._buffer_count - self._buffer_feats[0].shape[0] >= self.graph_size): + dropped = int(self._buffer_feats.pop(0).shape[0]) + del self._buffer_ids[:dropped] + self._buffer_count -= dropped + + def _graph_buffer(self): + """The most recent ``graph_size`` buffered samples, as one tensor. + + Rows that are not finite are dropped here -- a diverged batch must + never reach the encoder, whose weights would then turn NaN for good -- + and the buffer is rewritten to exactly what is returned. + """ + device = self._buffer_feats[-1].device + feats = th.cat([f.to(device) for f in self._buffer_feats], dim=0) + ids = list(self._buffer_ids) + if feats.shape[0] > self.graph_size: + feats = feats[-self.graph_size:] + ids = ids[-self.graph_size:] + finite = th.isfinite(feats).all(dim=1) + if not bool(finite.all()): + keep = finite.cpu().tolist() + dropped = keep.count(False) + self.skipped_nonfinite += dropped + if self.skipped_nonfinite == dropped: + logger.warning( + f"[projection] {dropped} samples have NaN/inf features (has the " + f"model diverged?); they are left out so the projection keeps " + f"its last good layout.") + feats = feats[finite] + ids = [i for i, k in zip(ids, keep) if k] + self._buffer_feats = [feats] if feats.shape[0] else [] + self._buffer_ids = ids + self._buffer_count = int(feats.shape[0]) + return feats, ids + + def _clear_buffer(self) -> None: + self._buffer_feats = [] + self._buffer_ids = [] + self._buffer_count = 0 + + def _fit(self, graph_feats: th.Tensor) -> None: + """A few encoder steps on the UMAP objective over *graph_feats*. + + Grad is switched back on here, whatever the caller had: per-sample + values are often logged from inside a ``no_grad`` (the Ultralytics + integration ships every signal under one), which left the UMAP loss + without a graph -- every fit failed, and after five the projection + turned itself off having never placed a point. The encoder and its + inputs are this tracker's own, so this cannot touch the model's graph. + """ + with th.inference_mode(False), th.enable_grad(): + if graph_feats.is_inference(): + # Captured under inference_mode: such a tensor cannot be + # saved for backward, a plain copy can. + graph_feats = graph_feats.clone() + target_p = membership_high_dim(graph_feats, self.n_neighbors) + loss = None + for _ in range(self.inner_steps): + self._optimizer.zero_grad(set_to_none=True) + loss = self._umap_loss(graph_feats, target_p) + if not bool(th.isfinite(loss)): + # Raised BEFORE the step, so a bad loss never reaches the + # encoder's weights. + raise FloatingPointError("non-finite UMAP loss") + loss.backward() + self._optimizer.step() + self.last_loss = float(loss.detach().item()) + self.steps_trained += 1 + + def observe_batch(self, batch_ids, step: int | None) -> bool: + """Pair the hooked features with *batch_ids*, fit a little, write the + coordinates. Returns True when this call wrote coordinates. + + Per training batch: buffer its features; place it with the current + encoder; and, every ``every_n_steps``, fit the encoder on the buffer and + re-place everything buffered. Per eval batch: place it, never fit. + + Every failure mode here is swallowed: this is instrumentation living + inside somebody's training loop, and it must not be able to take the + run down. + """ + if self._handle is None or batch_ids is None or self._writing_back: + return False + step_i = int(step or 0) + + with self._lock: + feats = self._pending + training = self._pending_training + if feats is None: + if step_i == self._last_step_seen: + # A second per-sample write for a batch already taken (a + # flag="loss" criterion AND a save_signals on the same + # batch). Nothing was missed. + return False + # The auto-picked layer never ran (an auxiliary head, a branch + # this forward skipped). Silence here means the board just never + # appears and nobody knows why, so say it once -- the fix is a + # one-line projection={"layer": ...}. + self._missed_batches += 1 + self._last_step_seen = step_i + if self._missed_batches == _MISS_WARN_AFTER: + logger.warning( + f"[projection] the hooked layer has not produced features " + f"in {_MISS_WARN_AFTER} visits -- it is probably not on this " + f"model's forward path. Point it somewhere else with " + f"wl.watch_or_edit(model, flag='model', " + f"projection={{'layer': ''}}), or " + f"set {ENV_ENABLED}=0 to turn the projection off." + ) + return False + self._pending = None + self._missed_batches = 0 + self._last_step_seen = step_i + try: + ids = [str(i) for i in ( + batch_ids.detach().cpu().tolist() + if hasattr(batch_ids, "detach") else list(batch_ids))] + except Exception: + return False + if len(ids) != feats.shape[0]: + # Mismatched pairing: writing it would attach one sample's + # features to another's id. A layer that never sees one row per + # sample (a per-point or per-token layer) does this on EVERY + # batch, and the board stayed empty without a word: say it once. + self.mismatched_batches += 1 + self.last_mismatch = (int(feats.shape[0]), len(ids)) + if self.mismatched_batches == _MISS_WARN_AFTER: + logger.warning( + f"[projection] the hooked layer {self.layer_name!r} gives " + f"{feats.shape[0]} feature rows for a batch of {len(ids)} " + f"samples -- not one row per sample (a per-point or " + f"per-token layer?), so nothing can be placed. Point it at " + f"a per-sample layer with wl.watch_or_edit(model, " + f"flag='model', projection={{'layer': ''}})." + ) + return False + + fit_now = False + if training: + self._push_buffer(feats, ids) + # The every_n_steps gate is about FITTING, and only training + # batches fit. (Applying it to eval batches too is why a whole + # split once had no coordinates: evaluation does not advance the + # step, so every batch of a validation pass carried the same one.) + fit_now = (step_i != self._last_step_run + and step_i % self.every_n_steps == 0) + if not fit_now: + # Place this batch with the encoder as it stands, so every + # sample of an epoch gets coordinates from a recent encoder + # -- not only the few batches that land on a fit step. Not + # before the first fit: an untrained encoder would draw a + # random cloud. + if self._encoder is None or self.steps_trained == 0: + return False + emit = feats + else: + # Only a fit claims the step. An eval batch that claimed it + # would lock out the training batch arriving at the same + # step, which is the very thing this guard exists for. + self._last_step_run = step_i + emit, ids = self._graph_buffer() + if emit.shape[0] < 4: + return False # too few samples for a kNN graph + else: + emit = feats + + try: + # Outside the caller's inference_mode, if any: an encoder + # built in there has weights autograd can never train. + with th.inference_mode(False): + self._ensure_encoder(emit.shape[1], emit.device) + if fit_now: + self._fit(emit) + # On a fit step, `emit` is the whole buffer: the encoder just + # moved, so every buffered sample's old coordinates are stale, + # and re-encoding them costs one more forward pass. Eval batches + # are placed, not learned from -- held-out samples belong in the + # picture, and letting them shape it would leak the eval set + # into it. An encoder that has never fit places them at random; + # the next fit overwrites those. + with th.no_grad(): + coords = self._encoder(emit).detach().cpu().numpy() + except Exception as exc: + self._record_failure(exc, step_i) + return False + if fit_now: + self._consecutive_failures = 0 + + # Never write a NaN over a sample's last good coordinates. + finite = np.isfinite(coords).all(axis=1) + if not finite.all(): + coords = coords[finite] + ids = [i for i, k in zip(ids, finite.tolist()) if k] + if not ids: + return False + return self._write_coords(ids, coords, step_i) + + def _write_coords(self, ids: list, coords: np.ndarray, step: int) -> bool: + """Persist coordinates as ordinary per-sample signals. + + Going through ``save_signals`` (rather than poking the dataframe) means + the projection inherits everything the rest of the pipeline already + has: H5 persistence, the metadata column list, sorting, histograms. + """ + try: + from weightslab.src import save_signals + axes = ["x", "y", "z"][:coords.shape[1]] + # Tensors, not numpy arrays: save_signals' own `normalize` handles + # torch tensors and lists, and returns None for a bare ndarray -- + # which lands the column in the dataframe holding nothing at all. + signals = { + f"{self.signal_prefix}_{axis}": + th.from_numpy(np.ascontiguousarray(coords[:, i], dtype=np.float32)) + for i, axis in enumerate(axes) + } + self._writing_back = True + try: + # _seen=False: placing a sample is not the model seeing it. + # Written as an ordinary step-tagged signal, every fit re-emits + # the whole buffer and so bumped nb_seen and last_seen on up to + # graph_size samples the model did not touch that step -- and + # counted every evaluated sample twice per pass. + save_signals(signals=signals, batch_ids=ids, step=step, log=False, + _seen=False) + finally: + self._writing_back = False + self.samples_written += len(ids) + if self._registered_prefix != self.signal_prefix: + # Tell the board these columns are a projection, and which + # layer they come from (see registry). + from weightslab.projection.registry import register_prefix + register_prefix(self.signal_prefix, layer=self.layer_name, + layer_detail=self.layer_detail) + self._registered_prefix = self.signal_prefix + # Checkpoint occasionally, so a crash or a Ctrl-C does not cost the + # layout. Cheap: a few hundred KB of encoder weights. + # + # Keyed on the fit that was last SAVED, not merely on the counter + # being divisible. `steps_trained` only advances on a fit, and this + # write-back now also runs for eval batches (they are placed, not + # learned from) -- so with the counter sitting on a multiple of + # _SAVE_EVERY_FITS, every batch of a full evaluation pass wrote the + # encoder to disk again. A single pass produced a burst of ~16 + # saves a second apart, for one fit's worth of weights. + if (self.steps_trained + and self.steps_trained % _SAVE_EVERY_FITS == 0 + and self.steps_trained != self._last_saved_fit): + self._last_saved_fit = self.steps_trained + save_projection(self) + return True + except Exception as exc: + if not getattr(self, "_warned_write", False): + self._warned_write = True + logger.warning( + f"[projection] could not write coordinates: {exc!r}. " + f"Training is unaffected; the board will not update.") + else: + logger.debug(f"[projection] coordinate write-back failed: {exc}") + return False + + def stats(self) -> dict: + return { + "enabled": self._handle is not None, + "out_dim": self.out_dim, + "every_n_steps": self.every_n_steps, + "graph_size": self.graph_size, + "buffered": len(self._buffer_ids), + "fits": self.steps_trained, + "rebuilds": self.rebuilds, + "feature_dim": self._encoder_dim, + "samples_written": self.samples_written, + "last_loss": self.last_loss, + "layer": self.layer_name, + "layer_detail": self.layer_detail, + "failures": self.failures, + "skipped_nonfinite": self.skipped_nonfinite, + "mismatched_batches": self.mismatched_batches, + "disabled_reason": self.disabled_reason, + "columns": [f"signals//{self.signal_prefix}_{a}" + for a in ["x", "y", "z"][:self.out_dim]], + } + + +# -------------------------------------------------------------------------- +# Module-level singleton: one projection per process, mirroring how the model, +# logger and dataframe are single registered objects in the ledger. +# -------------------------------------------------------------------------- +_TRACKER = None +# A checkpoint's encoder sidecar that was restored before any tracker existed +# (see restore_from_checkpoint); consumed by the next attach_projection. +_PENDING_CHECKPOINT = None + + +def get_tracker(): + return _TRACKER + + +# `aux_classifier` (torchvision FCN/DeepLab), `aux_head`/`auxiliary_head` +# (mmsegmentation), `aux` on its own. Matched against the dotted module name. +_AUXILIARY_NAME = re.compile(r"(^|\.)(aux|auxiliary)(_|\.|$)", re.IGNORECASE) + + +def pick_embedding_layer(model: nn.Module): + """Choose what to hook, for a model we know nothing about. + + The rule is one idea applied consistently: **hook the input of the last + parameterised layer**. That layer is the model's classifier/head, so what + flows INTO it is the representation -- the thing worth projecting -- while + what comes OUT is class scores. + + Getting this wrong is quiet and expensive. Reading a segmentation head's + *output* yields ``(B, num_classes)``, which is a perfectly well-shaped + tensor that UMAP will happily lay out; you would just be looking at a map + of class scores and believing it was the model's representation. + + The last parameterised layer is the last ``nn.Linear`` or conv in + registration order, whichever comes later: + + * an ``nn.Linear`` -> its input; features are on the LAST axis (Linear acts + on that axis), so ``(B, T, C)`` keeps ``C``; + * a conv (``Conv1d/2d/3d``, the head of a fully-convolutional model like + FCN/DeepLab/U-Net or a detector) -> its input; channels-first, so + ``(B, C, H, W)`` keeps ``C``; + * no parameterised layer at all -> the last leaf module's output, + channels-first, as a genuine last resort. + + "Whichever comes later", not "a Linear if there is one": a point-cloud + detector (PointPillars) has exactly one ``nn.Linear``, its per-point + feature net, registered FIRST, and a conv head. Preferring any Linear + hooked the per-point features -- one row per point, not per sample -- so + every batch was dropped and the board never got a point. + + AUXILIARY heads are skipped. torchvision's FCN and DeepLab register + ``aux_classifier`` AFTER ``classifier``, so "the last conv" picked the + auxiliary head -- which branches off an earlier backbone stage, so its + input is a shallower representation than the one the model actually + predicts from, and which torchvision does not even run outside training, so + on an evaluation pass the hook produced nothing at all and the cloud simply + stopped gaining test-split points. Both failures are silent. + + FROZEN layers are skipped too, while a trainable one exists. A layer whose + weights never train is a fixed operator, not the head: YOLO's ``Detect`` + registers its ``dfl`` box-decoding conv (frozen by construction) after the + class/box convs, so "the last conv" was the DFL -- which only runs to decode + boxes at inference, so the training forward never reached the hook, and + what it did see was a 16-bin box-distance softmax, not the representation. + A model frozen whole (wrapped for inspection only) still gets a pick. + + Returns ``(module, use_input, feature_last)``. + """ + # main = anything not under an auxiliary head; aux = the fallback, used + # only if a model turns out to be nothing BUT an auxiliary head; fixed = + # frozen parameterised layers, the last resort before a bare leaf. + last_param = {"main": None, "aux": None, "fixed": None} + last_leaf = {"main": None, "aux": None} + for name, module in model.named_modules(): + bucket = "aux" if _AUXILIARY_NAME.search(name) else "main" + param_bucket = "fixed" if _is_frozen(module) else bucket + if isinstance(module, (nn.Linear, nn.modules.conv._ConvNd)): + last_param[param_bucket] = module + if len(list(module.children())) == 0: + last_leaf[bucket] = module + + for bucket in ("main", "aux", "fixed"): + if last_param[bucket] is not None: + return _describe_layer(last_param[bucket]) + leaf = last_leaf["main"] or last_leaf["aux"] + return leaf, False, False + + +def _is_frozen(module: nn.Module) -> bool: + """Whether *module* has weights of its own and none of them train.""" + params = list(module.parameters(recurse=False)) + return bool(params) and not any(p.requires_grad for p in params) + + +def layer_detail(module: nn.Module, use_input: bool) -> str: + """What the hook reads at *module*, in words: ``"C3k2 output"``, + ``"Linear input"``. Shown beside the layer name on the projection board.""" + return f"{type(module).__name__} {'input' if use_input else 'output'}" + + +def _describe_layer(module): + """Hook policy for a chosen layer: same rule as the auto-pick -- read the + INPUT of a parameterised layer, and know which axis is which.""" + if isinstance(module, nn.Linear): + return module, True, True + if isinstance(module, nn.modules.conv._ConvNd): + return module, True, False + return module, False, False + + +def _resolve_layer(root: nn.Module, layer): + """``(module, use_input, feature_last, name)`` for *layer* on *root*. + + *layer* is a ``named_modules()`` name, a module, or ``None`` to auto-pick. + An unknown name falls back to the auto-pick with a warning: the live + projection must not stop a model from being wrapped. ``name`` is the + module's dotted name under *root* (``None`` if it is not under it), kept so + the same layer can be found again on a restored copy of the model. + """ + target, use_input, feature_last = None, True, True + if isinstance(layer, str): + found = dict(root.named_modules()).get(layer) + if found is None: + logger.warning(f"[projection] layer {layer!r} not found; auto-picking") + else: + target, use_input, feature_last = _describe_layer(found) + elif isinstance(layer, nn.Module): + target, use_input, feature_last = _describe_layer(layer) + if target is None: + target, use_input, feature_last = pick_embedding_layer(root) + if target is None: + return None, True, True, None + name = next((n for n, m in root.named_modules() if m is target), None) + return target, use_input, feature_last, name + + +def attach_projection(model: nn.Module, layer=None, **kwargs): + """Start projecting *model*'s features to 3-D. Returns the tracker, or + ``None`` when the feature is disabled or no layer could be hooked. + + *layer* may be a module, a ``named_modules()`` name, or ``None`` to + auto-pick (see :func:`pick_embedding_layer`). + + A previous tracker is detached first: there is one projection per process, + and a second ``watch_or_edit(..., flag="model")`` used to leave the first + model's hook installed, still filling a tracker nothing read. + """ + global _TRACKER, _PENDING_CHECKPOINT + if not projection_enabled(): + logger.debug(f"[projection] disabled via {ENV_ENABLED}") + _PENDING_CHECKPOINT = None + return None + + root = resolve_module(model) or model + target, use_input, feature_last, name = _resolve_layer(root, layer) + if target is None: + logger.debug("[projection] no hookable layer found") + return None + + kwargs.setdefault("out_dim", _env_int(ENV_DIM, 3)) + kwargs.setdefault("every_n_steps", _env_int(ENV_EVERY, 50)) + kwargs.setdefault("n_neighbors", _env_int(ENV_NEIGHBORS, 15)) + + detach_projection() + _purge_inert_hooks(root) + tracker = ProjectionTracker(**kwargs) + tracker.attach(target, use_input=use_input, feature_last=feature_last) + tracker.layer_name = name + tracker.layer_detail = layer_detail(target, use_input) + _TRACKER = tracker + + # Which encoder to start from. The model wrapper restores the checkpoint + # BEFORE the projection attaches, so when it restored one, the encoder + # saved WITH that checkpoint is the only one that matches the weights now + # in the model (see restore_from_checkpoint). Otherwise continue the run's + # own layout instead of re-laying it out from a random initialisation -- + # the coordinates already in the dataframe would describe a different + # embedding than the one about to overwrite them. + pending, _PENDING_CHECKPOINT = _PENDING_CHECKPOINT, None + if pending is not None: + _apply_checkpoint_encoder(tracker, pending) + else: + load_projection(tracker) + logger.info( + f"[projection] attached to {name or type(target).__name__} " + f"({tracker.layer_detail}, " + f"{'features last' if feature_last else 'channels first'}); " + f"disable with {ENV_ENABLED}=0 or projection=False" + ) + return tracker + + +def reattach_projection(model) -> bool: + """Point the live projection at *model*, keeping its encoder. + + For a checkpoint restore that registers a NEW model object (an + architecture restore unpickles one): the hook is still on a layer of the + old object, which nothing calls any more, so the projection would go quiet + with a misleading "layer not on the forward path" warning. The same layer + is looked up by name on the new model; the auto-pick covers a layer that + no longer exists there. + """ + tracker = _TRACKER + # A projection that turned itself off after repeated failures stays off. + if tracker is None or tracker.disabled_reason: + return False + root = resolve_module(model) + if root is None: + return False + _purge_inert_hooks(root) + target, use_input, feature_last, name = _resolve_layer(root, tracker.layer_name) + if target is None: + return False + tracker.attach(target, use_input=use_input, feature_last=feature_last) + tracker.layer_name = name + tracker.layer_detail = layer_detail(target, use_input) + tracker._pending = None + tracker._clear_buffer() + logger.info(f"[projection] re-attached to {name or type(target).__name__} " + f"on the restored model") + return True + + +def resolve_module(model): + """The real ``nn.Module`` behind whatever handle the caller is holding. + + ``wl.watch_or_edit(net, flag="model")`` hands back a ledger ``Proxy``, not + an ``nn.Module`` -- and that is the handle every training script keeps. On + it, ``named_modules()`` yields nothing, so a layer lookup against it finds + no layers and reports that none exist, which is a baffling thing to be told + about a model you can see. Descend the wrapper's ``.model`` chain to the + module that actually owns the layers. + + Returns ``None`` when there is no ``nn.Module`` to be found. + """ + try: + from weightslab.backend.model_interface import ModelInterface as _Wrapper + except Exception: # pragma: no cover - import cycle during package init + _Wrapper = () + current = model + for _ in range(6): # bounded: wrappers nest, but not deeply + if isinstance(current, nn.Module): + # A WeightsLab wrapper is itself a Module but delegates to .model; + # prefer the inner one, whose layer names are the ones the user + # sees in their own model definition. + # + # ONLY the WeightsLab wrapper. Any user module may hold a + # `self.model` (a backbone) beside its own head; descending into + # it would hook the backbone and never see the head's input. + inner = getattr(current, "model", None) if isinstance(current, _Wrapper) else None + if isinstance(inner, nn.Module) and inner is not current: + current = inner + continue + return current + # Not a Module yet (a Proxy, or a plain wrapper object): keep descending. + nxt = None + for attr in ("model", "module", "_model", "net"): + candidate = getattr(current, attr, None) + if candidate is not None and candidate is not current: + nxt = candidate + break + if nxt is None: + return None + current = nxt + return current if isinstance(current, nn.Module) else None + + +def _restart_pass(dataloader) -> bool: + """Make the next ``for`` over *dataloader* start at its first batch. + + A tracked WeightsLab loader is ONE stateful iterator: a ``for`` over it + carries on from wherever ``next(loader)`` left it (so a training loop and an + evaluation can share it). A training loop is nearly always mid-epoch, which + would hand a sweep only the rest of that epoch -- a few hundred samples out + of the split, with nothing to say so. Returns whether the loader could be + restarted; plain iterables (a torch ``DataLoader``) restart by themselves. + """ + restart = getattr(dataloader, "reset_iterator", None) + if not callable(restart): + return False + try: + restart() + except Exception as exc: + logger.debug(f"[projection] could not restart the loader's pass: {exc}") + return False + return True + + +@contextlib.contextmanager +def _frozen_age(model): + """Forward passes through *model* inside the block are not training steps. + + A tracked model counts a forward as a step whenever its tracking mode is + TRAIN -- and a training guard leaves it there on exit, so a sweep run + between two steps would silently age the model by one step per batch (the + plots' step axis jumps, and the next checkpoint is named after steps never + taken). The evaluation run does the same for the same reason: EVAL while it + runs, the previous mode back after. + """ + setter = getattr(model, "set_tracking_mode", None) + previous = None + if callable(setter): + try: + from weightslab.components.tracking import TrackingMode + previous = getattr(model, "tracking_mode", None) + setter(TrackingMode.EVAL) + except Exception as exc: + previous = None + logger.debug(f"[projection] could not freeze the model's age: {exc}") + try: + yield + finally: + if previous is not None: + try: + setter(previous) + except Exception: + pass + + +def _default_unpack(batch): + """``(inputs, sample_ids)`` from a WeightsLab batch. + + Every tracked loader yields ``(data, uid, ...)`` -- see the usecases, which + all destructure as ``for x, ids, y in loader``. Pass ``unpack=`` for a + loader that does something else. + """ + if isinstance(batch, (tuple, list)) and len(batch) >= 2: + return batch[0], batch[1] + raise ValueError( + "could not read (inputs, sample_ids) from the batch; pass " + "unpack=lambda batch: (inputs, ids)") + + +def project_dataset(model, dataloader, layer=None, prefix: str | None = None, + epochs: int = 10, fit_batch: int = 512, + max_samples: int | None = None, out_dim: int | None = None, + n_neighbors: int | None = None, lr: float = 1e-3, + min_dist: float = 0.1, spread: float = 1.0, + unpack=None, device=None, verbose: bool = True, + adopt: bool | None = None, method="umap") -> dict: + """Project a dataset **offline**, from a model that is not training. + + The live projection commits to one layer the moment training starts. This is + the answer to "that was the wrong layer": load the experiment as usual (the + checkpoint restores the weights), then re-project from whichever layer you + actually wanted, over the whole split, with no training step involved. + + It also fits *better* than the live path can. Online, the UMAP graph is + whatever one batch happened to contain; here every feature is collected + first, so each epoch draws its neighbourhoods from shuffled minibatches of + the **whole** dataset. + + *prefix* is the point of the exercise: writing to ``prefix="umap_layer3"`` + stores ``signals//umap_layer3_{x,y,z}`` alongside any existing projection + rather than overwriting it, so several candidate layers can coexist and be + switched between in the Studio board. + + *method* swaps the algorithm. ``"umap"`` (default) fits WeightsLab's own + parametric UMAP encoder. Anything else is YOUR projection, run on the + features collected here: an object with ``fit_transform`` (sklearn's + ``TSNE``/``PCA``, ``umap.UMAP``) or a callable ``(N, F) ndarray -> (N, 2|3)``. + It is stored through :func:`weightslab.save_projection_coords`, so no + encoder is saved or adopted, and ``epochs``/``fit_batch``/``lr``/ + ``min_dist``/``spread``/``n_neighbors``/``adopt`` do not apply. + + Args: + model: the model to read features from. Used in ``eval()``, under + ``no_grad`` -- its weights are never touched. + dataloader: any iterable of batches; by default ``batch[0]`` is the + input and ``batch[1]`` the sample ids (see :func:`_default_unpack`). + A tracked WeightsLab loader is restarted first, so the sweep covers + the whole split even when a training loop is mid-epoch on it (that + loop simply begins a fresh epoch afterwards). Discarded samples are + skipped by the loader itself, as in training. + layer: module, ``named_modules()`` name, or ``None`` to auto-pick. + prefix: signal prefix for the written coordinates. Defaults to + ``"umap"`` for the built-in method and to the method's name + (``TSNE(...)`` -> ``"tsne"``) for your own. + method: ``"umap"``, or your projection (see above). + epochs: passes over the collected features. + fit_batch: minibatch size for fitting. This is the UMAP graph size, so + bigger means better-connected neighbourhoods and more memory + (it is an N x N distance matrix). + max_samples: stop collecting after this many samples. + unpack: ``batch -> (inputs, sample_ids)`` override. + device: where to fit; defaults to the model's device. + adopt: hand the result to the LIVE projection, so that resuming training + continues refining *this* fit (same layer, same encoder, same + prefix) instead of starting a new one from scratch. + + ``None`` (default) means "adopt when it would otherwise corrupt": + writing to the prefix the live projection owns and NOT adopting + leaves training to overwrite these coordinates, sample by sample, + with a freshly-initialised encoder's output -- a cloud that is + part this layout and part another, which is meaningless. Pass + ``False`` to keep the live projection untouched anyway (fine when + *prefix* differs: that snapshot is then never written again). + + Returns: + dict: ``{"samples", "sample_ids", "feature_dim", "prefix", "columns", + "epochs", "final_loss", "layer", "method", "adopted_by_live"}``. + ``sample_ids`` names the samples placed, in the order of the feature + rows your *method* received -- so you can pair them with labels. + + Example: + >>> model = wl.watch_or_edit(MyNet(), flag="model") # weights restored + >>> loader = wl.watch_or_edit(train_ds, flag="data", loader_name="train_loader") + >>> wl.project_dataset(model, loader, layer="backbone.layer3", + ... prefix="umap_layer3", epochs=20) + >>> from sklearn.manifold import TSNE + >>> wl.project_dataset(model, loader, method=TSNE(n_components=3)) # -> 'tsne' + """ + from weightslab.projection import custom as _custom + + builtin = isinstance(method, str) and method.strip().lower() == "umap" + if isinstance(method, str) and not builtin: + raise ValueError( + f"project_dataset: method={method!r} is not built in. Pass 'umap', " + f"a callable features -> coords, or an object with fit_transform " + f"(e.g. sklearn.manifold.TSNE(n_components=3)).") + if builtin: + prefix = prefix or "umap" + else: + # Checked BEFORE the feature sweep, which can take minutes: a name the + # live projection owns would only be refused after all that work. + prefix = _custom.validate_prefix(prefix or _custom.method_name(method)) + _custom.check_not_live_prefix(prefix) + + unpack = unpack or _default_unpack + out_dim = out_dim if out_dim is not None else _env_int(ENV_DIM, 3) + n_neighbors = n_neighbors if n_neighbors is not None else _env_int(ENV_NEIGHBORS, 15) + + # The caller almost certainly holds the ledger Proxy that watch_or_edit + # returned, not the module itself. Hook and inspect the real module; keep + # calling the ORIGINAL handle for the forward pass, so whatever the wrapper + # does on the way through (masking, guards) still happens. + module = resolve_module(model) + if module is None: + raise TypeError( + "project_dataset: could not find an nn.Module on the object passed " + "as `model`. Pass the model (or the handle wl.watch_or_edit returned).") + + target, use_input, feature_last = None, True, True + if isinstance(layer, str): + found = dict(module.named_modules()).get(layer) + if found is None: + names = [n for n, _ in module.named_modules() if n] + raise ValueError( + f"project_dataset: layer {layer!r} not found. Available: " + f"{names[:20]}{'...' if len(names) > 20 else ''}") + target = found + elif isinstance(layer, nn.Module): + target = layer + if target is not None: + if isinstance(target, nn.Linear): + use_input, feature_last = True, True + elif isinstance(target, nn.modules.conv._ConvNd): + use_input, feature_last = True, False + else: + use_input, feature_last = False, False + else: + target, use_input, feature_last = pick_embedding_layer(module) + if target is None: + raise ValueError("project_dataset: no hookable layer found") + + tracker = ProjectionTracker(out_dim=out_dim, n_neighbors=n_neighbors, lr=lr, + min_dist=min_dist, spread=spread, + signal_prefix=prefix) + tracker.attach(target, use_input=use_input, feature_last=feature_last) + # Recorded with the prefix, so the board can say which layer this is. + tracker.layer_name = next((n for n, m in module.named_modules() if m is target), None) + tracker.layer_detail = layer_detail(target, use_input) + + was_training = module.training + module.eval() + feats_chunks, id_chunks = [], [] + collected = 0 + # The collection sweep is a full pass over the split with no output of its + # own -- on a real dataset that is minutes of silence, which reads as a + # hang. Same tqdm the ledger's own dataset passes use. + try: + total_batches = len(dataloader) + except TypeError: + total_batches = None + restarted = _restart_pass(dataloader) + stopped_early = False + sweep = tqdm(dataloader, total=total_batches, desc="[projection] features", + unit="batch", disable=not verbose, leave=False) + try: + with th.no_grad(), _frozen_age(model): + for batch in sweep: + inputs, ids = unpack(batch) + # Inputs follow the MODEL, always. `device` says where to fit + # the encoder, which is a separate question -- sending a cuda + # model's inputs to the cpu because the caller wanted to fit + # there just breaks the forward pass. + try: + inputs = inputs.to(next(module.parameters()).device) + except StopIteration: + pass + tracker._pending = None + model(inputs) + captured = tracker._pending + if captured is None: + continue + # Keep the collection on CPU: a whole split's features at + # once is exactly the thing that will not fit in VRAM. + feats_chunks.append(captured.cpu()) + id_chunks.extend( + [str(i) for i in (ids.detach().cpu().tolist() + if hasattr(ids, "detach") else list(ids))]) + collected += captured.shape[0] + if max_samples is not None and collected >= max_samples: + stopped_early = True + break + finally: + sweep.close() + tracker.detach() + if was_training: + module.train() + if restarted and stopped_early: + # Left alone, the training loop would carry on from where the sweep + # stopped, missing every batch it consumed. A finished sweep needs + # no such care: the exhausted loader begins a fresh epoch by itself. + _restart_pass(dataloader) + + if not feats_chunks: + raise RuntimeError( + f"project_dataset: {type(target).__name__} produced no features over " + f"the whole loader -- it is probably not on this model's forward path. " + f"Pick another with layer=.") + + feats = th.cat(feats_chunks, dim=0) + ids = id_chunks[:feats.shape[0]] + if max_samples is not None: + feats, ids = feats[:max_samples], ids[:max_samples] + if feats.shape[0] < 4: + raise RuntimeError( + f"project_dataset: only {feats.shape[0]} samples collected; need at " + f"least 4 to build a neighbourhood graph.") + + if not builtin: + if verbose: + logger.info( + f"[projection] {_custom.method_name(method)} on {feats.shape[0]} " + f"samples x {feats.shape[1]} features from {type(target).__name__} " + f"-> '{prefix}'") + coords = _custom.run_method(method, feats.numpy()) + columns = _custom.save_projection_coords(coords, ids, prefix=prefix) + from weightslab.projection.registry import register_prefix + register_prefix(prefix, layer=tracker.layer_name, + layer_detail=tracker.layer_detail) + return { + "samples": int(feats.shape[0]), + "sample_ids": ids, + "feature_dim": int(feats.shape[1]), + "prefix": prefix, + "columns": columns, + "epochs": 0, + "final_loss": None, + "layer": type(target).__name__, + "method": _custom.method_name(method), + "adopted_by_live": False, + } + + # Follow the model rather than grabbing the GPU because one exists: fitting + # on cuda for a cpu model leaves the encoder somewhere the live projection + # cannot use it if it later adopts this fit. + if device is not None: + fit_device = device + else: + try: + fit_device = next(module.parameters()).device + except StopIteration: + fit_device = feats.device + tracker._ensure_encoder(feats.shape[1], th.device(fit_device)) + if verbose: + logger.info( + f"[projection] offline fit: {feats.shape[0]} samples x {feats.shape[1]} " + f"features from {type(target).__name__}, {epochs} epochs -> '{prefix}'") + + n = feats.shape[0] + step = max(4, min(int(fit_batch), n)) + final_loss = None + epoch_bar = tqdm(range(max(1, int(epochs))), desc="[projection] fit", + unit="epoch", disable=not verbose, leave=False) + for epoch in epoch_bar: + order = th.randperm(n) + epoch_loss, batches = 0.0, 0 + for start in range(0, n - 3, step): + chunk = feats[order[start:start + step]].to(fit_device) + if chunk.shape[0] < 4: + continue + graph = membership_high_dim(chunk, n_neighbors) + tracker._optimizer.zero_grad(set_to_none=True) + loss = tracker._umap_loss(chunk, graph) + loss.backward() + tracker._optimizer.step() + epoch_loss += float(loss.detach().item()) + batches += 1 + final_loss = epoch_loss / max(1, batches) + tracker.steps_trained += batches + epoch_bar.set_postfix(loss=f"{final_loss:.4f}") + epoch_bar.close() + + coords = [] + with th.no_grad(): + for start in tqdm(range(0, n, step), desc="[projection] encode", + unit="chunk", disable=not verbose, leave=False): + coords.append(tracker._encoder(feats[start:start + step].to(fit_device)).cpu()) + coords = th.cat(coords, dim=0).numpy() + + tracker._write_coords(ids, coords, step=0) + tracker.last_loss = final_loss + save_projection(tracker) + + # --- hand over to the live projection, or say why we did not ----------- + live = get_tracker() + collides = live is not None and live.signal_prefix == prefix + should_adopt = collides if adopt is None else bool(adopt) + if should_adopt and live is not None: + # Re-point the live hook at the layer we just fitted on: adopting the + # encoder without the layer would feed it activations of a different + # width, which only triggers a rebuild and throws the fit away. + live.detach() + live.attach(target, use_input=use_input, feature_last=feature_last) + live.layer_name = tracker.layer_name + live.layer_detail = tracker.layer_detail + live._encoder = tracker._encoder + live._optimizer = tracker._optimizer + live._encoder_dim = tracker._encoder_dim + live._previous_dim = tracker._encoder_dim + live.signal_prefix = prefix + live.out_dim = tracker.out_dim + if verbose: + logger.info( + f"[projection] live projection adopted this fit: training will " + f"continue refining '{prefix}' from {type(target).__name__}") + elif collides: + logger.warning( + f"[projection] wrote prefix '{prefix}', which the LIVE projection also " + f"owns. Resuming training will overwrite these coordinates sample by " + f"sample with a freshly-initialised encoder, mixing two different " + f"layouts in one cloud. Use a distinct prefix= for a snapshot, or " + f"adopt=True to have training continue from this fit." + ) + + stats = { + "samples": int(n), + "sample_ids": ids, + "feature_dim": int(feats.shape[1]), + "prefix": prefix, + "columns": [f"signals//{prefix}_{a}" for a in ["x", "y", "z"][:out_dim]], + "epochs": int(epochs), + "final_loss": final_loss, + "layer": type(target).__name__, + "method": "umap", + "adopted_by_live": bool(should_adopt and live is not None), + } + if verbose: + logger.info(f"[projection] offline projection written: {stats['columns']}") + return stats + + +def checkpoint_path(root_log_dir=None, prefix: str = "umap"): + """Where a projection encoder is saved. + + Deliberately its OWN file, next to the run's logs, not inside the model + checkpoint: the encoder is not part of your model and must never ride along + into a weights file someone might deploy. ``None`` when no log directory + can be resolved. + """ + import os + from weightslab.projection.registry import projection_dir + folder = projection_dir(root_log_dir) + if folder is None: + return None + return os.path.join(folder, f"{prefix}_encoder.pt") + + +def save_projection(tracker=None, root_log_dir=None) -> str | None: + """Persist the encoder so a later run continues the SAME layout. + + Without this, restarting refits from a random initialisation and every + sample lands somewhere new -- the coordinates already in the dataframe then + describe a different embedding than the one still being written, which is + two layouts mixed in one cloud. + """ + import os + tracker = tracker or get_tracker() + state = encoder_state(tracker) if tracker is not None else None + if state is None: + return None + path = checkpoint_path(root_log_dir, tracker.signal_prefix) + if path is None: + return None + try: + os.makedirs(os.path.dirname(path), exist_ok=True) + th.save(state, path) + logger.info(f"[projection] encoder saved to {path}") + return path + except Exception as exc: + logger.debug(f"[projection] could not save encoder: {exc}") + return None + + +def encoder_state(tracker) -> dict | None: + """The tracker's encoder as a plain, ``th.save``-able dict (CPU tensors). + + One format for both places an encoder is written -- the run-level file and + the per-checkpoint sidecar -- so either can be read by + :func:`restore_encoder`. ``None`` before the first fit. + """ + if tracker is None or tracker._encoder is None: + return None + return { + "state_dict": {k: v.detach().cpu() for k, v in tracker._encoder.state_dict().items()}, + "in_dim": tracker._encoder_dim, + "out_dim": tracker.out_dim, + "prefix": tracker.signal_prefix, + "layer": tracker.layer_name, + "fits": tracker.steps_trained, + "a": tracker.a, + "b": tracker.b, + } + + +def restore_encoder(tracker, blob: dict) -> bool: + """Load an :func:`encoder_state` dict onto *tracker*. False on a mismatch. + + A different output dimension (``WEIGHTSLAB_PROJECTION_DIM`` changed since + the save) is a miss, not an error: the encoder could not produce the axes + the tracker writes. The buffer is cleared either way -- its features came + from whatever model was live before. + """ + try: + in_dim = int(blob.get("in_dim") or 0) + if in_dim <= 0 or int(blob.get("out_dim") or 0) != tracker.out_dim: + return False + encoder = ProjectionEncoder(in_dim, tracker.out_dim) + encoder.load_state_dict(blob["state_dict"]) + except Exception as exc: + logger.debug(f"[projection] unusable encoder state: {exc}") + return False + with tracker._lock: + tracker._encoder = encoder + tracker._encoder_dim = in_dim + tracker._previous_dim = in_dim + tracker._optimizer = th.optim.Adam(encoder.parameters(), lr=tracker.lr) + tracker.a = float(blob.get("a", tracker.a)) + tracker.b = float(blob.get("b", tracker.b)) + tracker.steps_trained = int(blob.get("fits", 0)) + tracker._last_saved_fit = tracker.steps_trained + tracker._pending = None + tracker._clear_buffer() + return True + + +def _apply_checkpoint_encoder(tracker, sidecar) -> bool: + """Set *tracker* to the state saved beside a checkpoint. + + No sidecar means no encoder existed when that checkpoint was taken (it is + written with every checkpoint once there is one), so the tracker is RESET + rather than left holding a later step's layout: restoring the model to + step 200 and keeping the step-5000 encoder would read old features through + a newer map and draw a layout that never existed. + """ + import os + path = str(sidecar) + if os.path.exists(path): + try: + blob = th.load(path, map_location="cpu", weights_only=False) + except Exception as exc: + logger.warning(f"[projection] could not read encoder {path}: {exc!r}; " + f"the projection restarts from scratch") + blob = None + if blob is not None and restore_encoder(tracker, blob): + logger.info( + f"[projection] encoder restored with the checkpoint " + f"({tracker._encoder_dim} -> {tracker.out_dim}, " + f"{tracker.steps_trained} prior fits)") + return True + tracker.reset() + logger.info("[projection] the restored checkpoint predates the projection's " + "first fit; the layout restarts from that point") + return False + + +def restore_from_checkpoint(sidecar) -> bool: + """Bring the live projection to the moment a model checkpoint was taken. + + Called by the checkpoint manager whenever it restores weights, with the + path of that checkpoint's encoder sidecar (which may not exist). + + There are two orders this can happen in. A restore from the Studio or the + agent finds the tracker attached and applies at once. The restore at + START-UP does not: ``ModelInterface`` reloads the latest checkpoint while + it is being constructed, and ``watch_or_edit`` only attaches the + projection after that. Applying "now" there was a silent no-op -- the + tracker then attached and picked up the run-level encoder, a later one + than the weights just restored. So with no tracker the sidecar is parked, + and the next :func:`attach_projection` starts from it. + + Returns True when an encoder was (or, parked, will be) restored. + """ + global _PENDING_CHECKPOINT + import os + tracker = _TRACKER + if tracker is None: + _PENDING_CHECKPOINT = str(sidecar) + return os.path.exists(str(sidecar)) + return _apply_checkpoint_encoder(tracker, sidecar) + + +def load_projection(tracker=None, root_log_dir=None) -> bool: + """Restore a saved encoder onto *tracker*. True when one was loaded. + + A width mismatch (the hooked layer changed since the save) is a miss, not + an error: the tracker simply starts fresh rather than loading an encoder + that cannot accept the features it will be given. + """ + import os + tracker = tracker or get_tracker() + if tracker is None: + return False + path = checkpoint_path(root_log_dir, tracker.signal_prefix) + if not path or not os.path.exists(path): + return False + try: + blob = th.load(path, map_location="cpu", weights_only=False) + except Exception as exc: + logger.debug(f"[projection] could not load encoder: {exc}") + return False + if not restore_encoder(tracker, blob): + return False + logger.info( + f"[projection] encoder restored from {path} " + f"({tracker._encoder_dim} -> {tracker.out_dim}, {tracker.steps_trained} prior fits); " + f"the layout continues rather than restarting") + return True + + +def mirror_projection(model) -> bool: + """Feed the live projection from *model* too: a second copy of the + projected model that was made BEFORE the projection attached. + + Ultralytics is the case in point: its trainer builds the EMA shadow in + ``_setup_train``, before ``on_train_start`` where the model gets wrapped, + and validation runs on that shadow -- so without this no validation sample + is ever placed. A copy made AFTER the attach needs nothing; deepcopy keeps + the hook. The same layer is found by name; returns whether *model* now + feeds the tracker. + """ + tracker = _TRACKER + if tracker is None or tracker.disabled_reason or not tracker.layer_name: + return False + root = resolve_module(model) + if root is None: + return False + target = dict(root.named_modules()).get(tracker.layer_name) + if target is None: + logger.debug(f"[projection] no layer {tracker.layer_name!r} to mirror on " + f"{type(root).__name__}") + return False + return tracker.mirror(target) + + +def detach_projection() -> None: + """Remove the hook and drop the tracker.""" + global _TRACKER + if _TRACKER is not None: + _TRACKER.detach() + _TRACKER = None + + +def clear_pending_restore() -> None: + """Forget a parked checkpoint restore (see :func:`restore_from_checkpoint`).""" + global _PENDING_CHECKPOINT + _PENDING_CHECKPOINT = None + + +def observe_batch(batch_ids, step) -> bool: + """Called from the loss wrapper for every logged batch. No-op when the + feature is off, which is the common case for the hot path.""" + tracker = _TRACKER + if tracker is None: + return False + try: + return tracker.observe_batch(batch_ids, step) + except Exception as exc: + logger.debug(f"[projection] observe_batch failed: {exc}") + return False diff --git a/weightslab/projection/registry.py b/weightslab/projection/registry.py new file mode 100644 index 00000000..63949ea1 --- /dev/null +++ b/weightslab/projection/registry.py @@ -0,0 +1,196 @@ +"""Which column prefixes are projections. + +The board used to offer every prefix that had both an ``_x`` and a ``_y`` +column, so an ordinary pair of signals -- ``center_x`` / ``center_y`` on a +detection set, ``pos_x`` / ``pos_y`` on a tabular one -- showed up in the picker +as a "projection" called ``center``. Everything that writes a projection goes +through WeightsLab (the live tracker, ``project_dataset``, +``save_projection_coords``), so each writer records its prefix here instead, +and the board offers exactly those. + +Kept in memory and in ``/projection/prefixes.json``, so a +restarted run still knows which of its columns are projections. A run with no +registry at all (written before it existed) falls back to the column pairs. + +Each prefix also records the layer its features were read from, when +WeightsLab computed it (the board shows it under its title). The file is a +JSON object ``{prefix: {"layer": ..., "layer_detail": ...}}``; the plain list +older versions wrote still reads, and an older version reading the object +iterates its keys -- the same prefixes. +""" + +from __future__ import annotations + +import json +import logging +import os +import threading + +logger = logging.getLogger(__name__) + +REGISTRY_FILE = "prefixes.json" + +_LOCK = threading.Lock() +_REGISTERED: set = set() +# prefix -> {"layer": ..., "layer_detail": ...}, for the prefixes registered here. +_INFO: dict = {} +# path -> ((mtime, size), {prefix: info}), so the board's polling does not re-read the +# file. The size is part of the key because mtime alone is not enough: two writes +# inside one filesystem timestamp tick (coarse on Windows) share an mtime, and +# the second read would return the first write's prefixes. +_FILE_CACHE: dict = {} + + +def projection_dir(root_log_dir=None): + """``/projection``, or ``None`` when no root can be resolved. + + Falsy, not just None, at every step: an unset hyperparameter comes back as + "" rather than None, and testing ``is None`` there skipped the env fallback. + """ + root = root_log_dir + if not root: + try: + from weightslab.backend.ledgers import get_hyperparams + root = get_hyperparams().get("root_log_dir") + except Exception: + root = None + if not root: + root = os.environ.get("WEIGHTSLAB_ROOT_LOG_DIR") + if not root: + return None + return os.path.join(str(root), "projection") + + +def _registry_path(root_log_dir=None): + folder = projection_dir(root_log_dir) + return os.path.join(folder, REGISTRY_FILE) if folder else None + + +def _read_file(path) -> dict: + """``{prefix: info}`` from the registry file; ``{}`` when unreadable.""" + try: + stamp = (os.path.getmtime(path), os.path.getsize(path)) + except OSError: + return {} + cached = _FILE_CACHE.get(path) + if cached and cached[0] == stamp: + return dict(cached[1]) + try: + with open(path, "r", encoding="utf-8") as handle: + raw = json.load(handle) + except Exception as exc: + logger.debug(f"[projection] unreadable registry {path}: {exc}") + return {} + items = raw.items() if isinstance(raw, dict) else ((n, {}) for n in raw) + entries = {str(n): (dict(info) if isinstance(info, dict) else {}) + for n, info in items if str(n).strip()} + _FILE_CACHE[path] = (stamp, entries) + return dict(entries) + + +def _write_file(path, entries: dict) -> None: + os.makedirs(os.path.dirname(path), exist_ok=True) + tmp = path + ".tmp" + with open(tmp, "w", encoding="utf-8") as handle: + json.dump({name: entries[name] for name in sorted(entries)}, handle) + os.replace(tmp, path) + + +def register_prefix(prefix: str, root_log_dir=None, layer=None, + layer_detail=None) -> None: + """Record *prefix* as a projection. Best effort: never raises. + + *layer* (a ``named_modules()`` name) and *layer_detail* (e.g. ``"C3k2 + output"``) say where its features came from. Left out, whatever was + recorded before is kept: a projection the caller computed itself + (:func:`save_projection_coords`) has no layer WeightsLab knows of. + """ + prefix = str(prefix).strip() + if not prefix: + return + info = {key: str(value) for key, value in + (("layer", layer), ("layer_detail", layer_detail)) if value} + with _LOCK: + new = prefix not in _REGISTERED + _REGISTERED.add(prefix) + if info: + _INFO[prefix] = {**_INFO.get(prefix, {}), **info} + path = _registry_path(root_log_dir) + if path is None: + return + on_disk = _read_file(path) + merged = {**on_disk.get(prefix, {}), **info} + if not new and prefix in on_disk and merged == on_disk[prefix]: + return + on_disk[prefix] = merged + try: + _write_file(path, on_disk) + except Exception as exc: + logger.debug(f"[projection] could not persist prefix {prefix!r}: {exc}") + + +def prefix_info(prefix: str, root_log_dir=None) -> dict: + """What is recorded about *prefix*: ``{"layer": ..., "layer_detail": ...}``, + either key absent when unknown. Never raises.""" + prefix = str(prefix).strip() + with _LOCK: + info = dict(_INFO.get(prefix, {})) + try: + path = _registry_path(root_log_dir) + except Exception: + path = None + if path is not None and os.path.exists(path): + # This process's own record wins; the file is what a restarted + # run knows. + info = {**_read_file(path).get(prefix, {}), **info} + return info + + +def unregister_prefix(prefix: str, root_log_dir=None) -> None: + """Stop offering *prefix* as a projection. Best effort: never raises. + + The columns are left where they are; this only takes the name out of the + board's picker. Used when a checkpoint restore returns the run to a moment + that predates the projection. + """ + prefix = str(prefix).strip() + if not prefix: + return + with _LOCK: + _REGISTERED.discard(prefix) + _INFO.pop(prefix, None) + path = _registry_path(root_log_dir) + if path is None or not os.path.exists(path): + return + on_disk = _read_file(path) + if prefix not in on_disk: + return + on_disk.pop(prefix) + try: + _write_file(path, on_disk) + except Exception as exc: + logger.debug(f"[projection] could not drop prefix {prefix!r}: {exc}") + + +def known_prefixes(root_log_dir=None): + """Every registered prefix, or ``None`` when there is no registry at all. + + ``None`` (not an empty set) is what tells the caller to fall back to + discovering prefixes from the columns, for runs written before this + registry existed. + """ + with _LOCK: + names = set(_REGISTERED) + path = _registry_path(root_log_dir) + if path is not None and os.path.exists(path): + names |= set(_read_file(path)) + return names + return names or None + + +def clear_registry() -> None: + """Forget the in-memory registry (the file is left alone).""" + with _LOCK: + _REGISTERED.clear() + _INFO.clear() + _FILE_CACHE.clear() diff --git a/weightslab/projection/snapshots.py b/weightslab/projection/snapshots.py new file mode 100644 index 00000000..2b0d326e --- /dev/null +++ b/weightslab/projection/snapshots.py @@ -0,0 +1,184 @@ +"""Your own projections, kept with the checkpoint they were computed at. + +The built-in UMAP encoder is saved beside every model checkpoint and restored +with it (see ``CheckpointManager._save_projection_sidecar``). A projection you +computed yourself (``wl.save_projection_coords`` / ``wl.project_dataset(..., +method=...)``) is just per-sample columns, which are not part of the model -- +so on its own it would describe whichever model state is current *now*, however +far a restore had since moved the weights. + +This module closes that gap. At every checkpoint the coordinates of each custom +projection are written to a sidecar next to the weights; restoring that +checkpoint puts exactly those projections back and takes the others out of the +picker. "The projections I can see" is then always the set that belongs to the +model I am looking at. + +Nothing here is user-facing: the checkpoint manager calls it. +""" + +from __future__ import annotations + +import logging + +import numpy as np +import pandas as pd + +from weightslab.projection.custom import AXES +from weightslab.projection.registry import ( + known_prefixes, register_prefix, unregister_prefix) + +logger = logging.getLogger(__name__) + +SIDECAR_SUFFIX = ".coords.pt" +SIGNAL_PREFIX = "signals//" + + +def sidecar_path(checkpoint_file): + """``projection/.coords.pt`` beside *checkpoint_file*. + + The ``projection/`` subdirectory for the same reason the encoder's sidecar + lives there: the weight-checkpoint glob must never see it. + """ + from pathlib import Path + path = Path(checkpoint_file) + return path.parent / "projection" / (path.stem + SIDECAR_SUFFIX) + + +def custom_prefixes() -> list: + """Registered projections that are the user's own, sorted. + + The live encoder's prefix is excluded: it is snapshotted as an encoder, and + re-placing its samples is the live projection's job. + """ + names = known_prefixes() + if not names: + return [] + try: + from weightslab.projection.parametric_umap import get_tracker + live = get_tracker() + live_prefix = live.signal_prefix if live is not None else None + except Exception: + live_prefix = None + return sorted(n for n in names if n != live_prefix) + + +def _sample_rows(frame: pd.DataFrame) -> pd.DataFrame: + """Sample-level rows only (``annotation_id == 0``), indexed by sample id.""" + names = list(getattr(frame.index, "names", []) or []) + if "annotation_id" in names: + frame = frame.xs(0, level="annotation_id") + return frame + + +def snapshot(frame: pd.DataFrame, prefixes=None) -> dict: + """``{prefix: {"ids": [str], "coords": float32 (n, 2|3)}}`` from *frame*. + + Only samples with a finite position are kept (NaN means "not placed"), and a + prefix with no complete column pair is skipped. + """ + out = {} + if frame is None or len(frame) == 0: + return out + rows = _sample_rows(frame) + for prefix in (custom_prefixes() if prefixes is None else prefixes): + columns = [] + for axis in AXES: + name = f"{SIGNAL_PREFIX}{prefix}_{axis}" + if name not in rows.columns: + break + columns.append(name) + if len(columns) < 2: + continue + coords = np.column_stack([ + pd.to_numeric(rows[c], errors="coerce").to_numpy(dtype=np.float32) + for c in columns]) + finite = np.isfinite(coords).all(axis=1) + if not finite.any(): + continue + ids = rows.index[finite] + out[prefix] = {"ids": [str(i) for i in ids], "coords": coords[finite]} + return out + + +def save_sidecar(checkpoint_file, frame: pd.DataFrame) -> bool: + """Write the custom projections held by *frame* beside *checkpoint_file*. + + Written even when there are none: an empty snapshot is what says "this + moment had no custom projection", so restoring it clears later ones. + """ + import torch as th + target = sidecar_path(checkpoint_file) + target.parent.mkdir(parents=True, exist_ok=True) + th.save({"version": 1, "projections": snapshot(frame)}, target) + return True + + +def load_sidecar(checkpoint_file): + """The saved ``{prefix: {...}}`` for *checkpoint_file*, or ``None`` when the + checkpoint predates this feature (so a restore leaves things alone).""" + import torch as th + path = sidecar_path(checkpoint_file) + if not path.exists(): + return None + payload = th.load(path, weights_only=False, map_location="cpu") + return dict(payload.get("projections", {})) + + +def apply(projections: dict, frame: pd.DataFrame) -> dict: + """Make the run's custom projections exactly *projections*. + + * each saved projection is written back (all of its samples); + * every other custom projection is blanked to NaN and dropped from the + picker -- it was computed after this moment, so it does not exist yet. + + Returns ``{"restored": [...], "cleared": [...]}``. + """ + from weightslab import src as _src + + saved = set(projections) + present = set(custom_prefixes()) + cleared = [] + + rows = _sample_rows(frame) if frame is not None and len(frame) else None + all_ids = [str(i) for i in rows.index] if rows is not None else [] + + for prefix in sorted(present - saved): + axes = [a for a in AXES + if rows is not None and f"{SIGNAL_PREFIX}{prefix}_{a}" in rows.columns] + if all_ids and axes: + nan = np.full(len(all_ids), np.nan, dtype=np.float32) + _src.save_signals( + signals={f"{prefix}_{a}": nan for a in axes}, + batch_ids=all_ids, log=False, _seen=False) + unregister_prefix(prefix) + cleared.append(prefix) + + restored = [] + for prefix in sorted(saved): + entry = projections[prefix] + coords = np.asarray(entry["coords"], dtype=np.float32) + ids = [str(i) for i in entry["ids"]] + if not ids or coords.ndim != 2: + continue + signals = {f"{prefix}_{a}": coords[:, i] + for i, a in enumerate(AXES[:coords.shape[1]])} + # Samples the snapshot did not place hold whatever a later write left: + # blank them first so the cloud is the saved one, not a blend. + if all_ids: + blank = np.full(len(all_ids), np.nan, dtype=np.float32) + _src.save_signals(signals={k: blank for k in signals}, + batch_ids=all_ids, log=False, _seen=False) + _src.save_signals(signals=signals, batch_ids=ids, log=False, _seen=False) + register_prefix(prefix) + restored.append(prefix) + + try: + frame_m = _src.DATAFRAME_M if _src.DATAFRAME_M is not None else _src.get_dataframe() + if hasattr(frame_m, "flush"): + frame_m.flush() + except Exception as exc: + logger.debug(f"[projection] snapshot restore: flush failed ({exc})") + if restored or cleared: + logger.info(f"[projection] checkpoint restore: custom projections " + f"restored={restored} cleared={cleared}") + return {"restored": restored, "cleared": cleared} diff --git a/weightslab/proto/experiment_service.proto b/weightslab/proto/experiment_service.proto index 17386eb8..7e69c5bd 100644 --- a/weightslab/proto/experiment_service.proto +++ b/weightslab/proto/experiment_service.proto @@ -98,6 +98,12 @@ service ExperimentService { // (or a zip, for formats that need one file per image) comes back as bytes // in a single response. rpc ExportAnnotations (ExportAnnotationsRequest) returns (ExportAnnotationsResponse); + // The live parametric-UMAP projection of the model's representation, as a + // 3-D point per sample. Level-of-detail is server-side and driven by the + // client's current view box: narrowing the box as the user zooms in buys + // finer detail for the SAME point budget, instead of ever shipping the whole + // dataset. See ProjectionRequest. + rpc GetProjection (ProjectionRequest) returns (ProjectionResponse); } // --- Logger Data Sync --- @@ -755,6 +761,10 @@ message DataSplitsResponse { message AgentHealthResponse { bool available = 1; string message = 2; + // The model the NEXT query will use, re-resolved from OpenCode's config. + // Health is polled, so this is how the studio notices a model chosen from + // the CLI (or by another client) instead of only knowing what it set itself. + string current_model = 3; } // --- Agent Onboarding --- @@ -794,6 +804,9 @@ message GetAgentModelsResponse { bool success = 1; repeated string models = 2; string message = 3; + // Which of `models` is active right now, so the picker can preselect it + // rather than opening on whatever the client last set. + string current_model = 4; } message ResetAgentResponse { @@ -989,3 +1002,139 @@ message ExportAnnotationsResponse { string mime_type = 5; // e.g. "application/xml", "application/json", "application/zip" int32 image_count = 6; // number of images included, for UI feedback } + +// --- Live 3-D projection (parametric UMAP) --- + +// A request for the points visible in one view box. +// +// The whole level-of-detail scheme lives in the pairing of `bounds` with +// `max_points`: the client sends the box its camera currently frames and a +// budget it can actually draw, and the server decimates what falls inside that +// box down to the budget. Zooming in shrinks the box, so the same budget buys +// progressively finer detail -- which is why a multi-million-sample projection +// never needs to cross the wire. +message ProjectionRequest { + bool has_bounds = 1; // false => whole extent (the initial, fully zoomed-out view) + float min_x = 2; + float min_y = 3; + float min_z = 4; + float max_x = 5; + float max_y = 6; + float max_z = 7; + int32 max_points = 8; // client render budget; 0 => server default + string color_column = 9; // optional dataframe column to colour points by ("" = by split) + string prefix = 10; // coordinate signal prefix; "" => "umap" + + // Follow the data board's current filter. + // + // True for an ordinary refresh: a filter the user (or the agent) applied to + // the data view should narrow the projection with it. False only when the + // projection ITSELF produced the active filter -- a lasso pushes a sample_id + // query into the data board, and re-reading it here would collapse the cloud + // to the points just selected, leaving nothing to select from next. + bool follow_view = 11; + + // The samples the data grid is rendering right now (its current page). + // + // Honoured whenever follow_view is set: the cloud narrows to exactly these + // samples, filtered data view or not. An empty list means "no restriction". + // Served in full: no view box, no budget -- a grid page is small. + repeated string restrict_sample_ids = 12; + + // The camera's view frustum, as 6 planes of 4 floats each (nx, ny, nz, d): + // a point p is on screen when n.p + d >= 0 for every plane. + // + // Sent alongside the min/max box, which is then the frustum's bounding box + // clipped to the cloud. The box alone is a poor stand-in for a perspective + // view: it either cuts off points that are on screen or spends the budget on + // points that are not. Empty => the box is the whole region. + repeated float frustum_planes = 13; + + // --- Lasso selection ------------------------------------------------------ + // With a mode set, the request is a SELECTION rather than a view: every + // sample whose screen position falls inside `lasso_px` -- drawn or not, since + // the view only ever draws a subsample -- is selected on the server and kept + // there under ProjectionResponse.selection_token, because a selection over + // millions of samples cannot travel as a list of ids. The box and planes + // above then describe the lasso's own sub-frustum and only narrow the + // candidates. + ProjectionLassoMode lasso_mode = 14; + repeated float lasso_px = 15; // polygon, x,y pairs in canvas pixels + repeated float lasso_viewport = 16; // canvas width, height in pixels + repeated float view_projection = 17; // camera projection * view, column-major 4x4 +} + +enum ProjectionLassoMode { + PROJECTION_LASSO_NONE = 0; // an ordinary view request + PROJECTION_LASSO_REPLACE = 1; // the selection becomes what is inside the loop + PROJECTION_LASSO_ADD = 2; // ... plus what was selected before + // ... minus what is outside it: only what was selected before AND is inside + // the loop. Drawn from another angle, this carves a 3-D volume out of the + // column a single loop selects. + PROJECTION_LASSO_REFINE = 3; +} + +// Coordinates come back as PACKED parallel arrays rather than one message per +// point: at a 50k budget the per-message framing of a `repeated ProjectionPoint` +// costs more than the coordinates themselves, and Float32Array/typed-array +// upload into a three.js BufferGeometry is exactly the shape the renderer wants. +message ProjectionResponse { + bool success = 1; + string message = 2; + + repeated string sample_ids = 3; // length N; identifies each point for selection + repeated float coords = 4; // length 3N: x,y,z triples, point-major + repeated int32 origin_ids = 5; // length N; index into `origins` + repeated string origins = 6; // split names, for colouring + repeated float color_values = 7; // length N when color_column was numeric; else empty + repeated string color_labels = 8; // length N when color_column was categorical; else empty + + int64 returned = 9; // N + int64 total_in_view = 10; // rows inside the box BEFORE decimation + int64 total_available = 11; // rows that have coordinates at all + + // Full extent of the projection (NOT of the requested box), so the client can + // frame the whole cloud on first load and show how far in the current view is. + float extent_min_x = 12; + float extent_min_y = 13; + float extent_min_z = 14; + float extent_max_x = 15; + float extent_max_y = 16; + float extent_max_z = 17; + + int32 fits = 18; // encoder fits so far, so the UI can say how fresh this is + int32 dims = 19; // 2 or 3 + + // Every projection stored on this dataset, by prefix ("umap", plus anything + // wl.project_dataset() wrote under its own name). Offline re-projections are + // the whole reason this exists: having re-projected from a different layer, + // the user needs to be able to SWITCH to it, and only the server knows which + // ones are present. Feed one back as ProjectionRequest.prefix. + repeated string available_prefixes = 20; + + // Per-point discarded flag (length N). Discarded samples are drawn grey + // rather than in their class/split colour: they are still part of the + // representation and worth seeing in place, but they are not training the + // model, so they must not read as live members of a cluster. + repeated bool discarded = 21; + + // Lasso requests only: how many samples the selection holds, and the token + // that applies it to the data board -- send the data query + // "@projection_selection ". The ids themselves stay on the server. + int64 selected = 22; + string selection_token = 23; + + // View requests: per point (length N), whether it is in the current lasso + // selection. Sent so a point that only arrives with a later zoom is drawn + // selected too -- the client never holds the whole selection. Empty when + // nothing is selected. + repeated bool in_selection = 24; + + // Which layer the coordinates were computed from, so the board can say what + // it is a picture of: its named_modules() name ("model.22" -- what + // projection={'layer': ...} and project_dataset(layer=...) take) and what + // was read there ("C3k2 output", "Linear input"). Empty for a projection + // WeightsLab did not compute (wl.save_projection_coords). + string layer = 25; + string layer_detail = 26; +} diff --git a/weightslab/proto/experiment_service_pb2.py b/weightslab/proto/experiment_service_pb2.py index 3709cc25..5137d3b9 100644 --- a/weightslab/proto/experiment_service_pb2.py +++ b/weightslab/proto/experiment_service_pb2.py @@ -24,7 +24,7 @@ -DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile(b'\n)weightslab/proto/experiment_service.proto\"\x81\x02\n\x1aGetLatestLoggerDataRequest\x12\x1c\n\x14request_full_history\x18\x01 \x01(\x08\x12\x12\n\nmax_points\x18\x02 \x01(\x05\x12\x17\n\x0f\x62reak_by_slices\x18\x03 \x01(\x08\x12\x0c\n\x04tags\x18\x04 \x03(\t\x12\x12\n\ngraph_name\x18\x05 \x01(\t\x12\r\n\x05x_min\x18\x06 \x01(\x03\x12\r\n\x05x_max\x18\x07 \x01(\x03\x12\x13\n\x0bhas_x_range\x18\x08 \x01(\x08\x12\x14\n\x0cmetric_names\x18\t \x03(\t\x12\x19\n\x11\x65xperiment_hashes\x18\n \x03(\t\x12\x12\n\nindex_only\x18\x0b \x01(\x08\"\xa2\x01\n\x10SignalCurveIndex\x12\x13\n\x0bmetric_name\x18\x01 \x01(\t\x12\x17\n\x0f\x65xperiment_hash\x18\x02 \x01(\t\x12\x12\n\nfirst_step\x18\x03 \x01(\x03\x12\x11\n\tlast_step\x18\x04 \x01(\x03\x12\x13\n\x0bpoint_count\x18\x05 \x01(\x03\x12\x11\n\tvalue_min\x18\x06 \x01(\x01\x12\x11\n\tvalue_max\x18\x07 \x01(\x01\"1\n\rSignalOutlier\x12\x11\n\tsample_id\x18\x01 \x01(\t\x12\r\n\x05value\x18\x02 \x01(\x02\"\xd2\x03\n\x0fLoggerDataPoint\x12\x13\n\x0bmetric_name\x18\x01 \x01(\t\x12\x11\n\tmodel_age\x18\x02 \x01(\x05\x12\x14\n\x0cmetric_value\x18\x03 \x01(\x02\x12\x17\n\x0f\x65xperiment_hash\x18\x04 \x01(\t\x12\x11\n\ttimestamp\x18\x05 \x01(\x03\x12\x11\n\tsample_id\x18\x06 \x01(\t\x12\x1c\n\x14is_evaluation_marker\x18\x07 \x01(\x08\x12\x12\n\nsplit_name\x18\x08 \x01(\t\x12\x17\n\x0f\x65valuation_tags\x18\t \x03(\t\x12\x12\n\npoint_note\x18\n \x01(\t\x12\x12\n\naudit_mode\x18\x0b \x01(\x08\x12 \n\x08outliers\x18\x0c \x03(\x0b\x32\x0e.SignalOutlier\x12\x15\n\routlier_count\x18\r \x01(\x05\x12\x14\n\x0csample_count\x18\x0e \x01(\x05\x12\x13\n\x0btrend_value\x18\x0f \x01(\x02\x12\x14\n\x0ctrend_margin\x18\x10 \x01(\x02\x12\x16\n\x0ehas_trend_band\x18\x11 \x01(\x08\x12\x11\n\tvalue_min\x18\x12 \x01(\x02\x12\x11\n\tvalue_max\x18\x13 \x01(\x02\x12\x17\n\x0fhas_value_range\x18\x14 \x01(\x08\"\x9a\x01\n\x1bGetLatestLoggerDataResponse\x12 \n\x06points\x18\x01 \x03(\x0b\x32\x10.LoggerDataPoint\x12\x1a\n\x12weightslab_version\x18\x02 \x01(\t\x12!\n\x06\x63urves\x18\x03 \x03(\x0b\x32\x11.SignalCurveIndex\x12\x1a\n\x12\x61pplied_max_points\x18\x04 \x01(\x05\"\x07\n\x05\x45mpty\"/\n\x08NeuronId\x12\x10\n\x08layer_id\x18\x01 \x01(\x05\x12\x11\n\tneuron_id\x18\x02 \x01(\x05\"\x91\x02\n\x0fWeightOperation\x12*\n\x07op_type\x18\x01 \x01(\x0e\x32\x14.WeightOperationTypeH\x00\x88\x01\x01\x12\x15\n\x08layer_id\x18\x02 \x01(\x05H\x01\x88\x01\x01\x12\x1d\n\nneuron_ids\x18\x03 \x03(\x0b\x32\t.NeuronId\x12\x16\n\x0eneurons_to_add\x18\t \x01(\x05\x12 \n\x18zerofy_from_incoming_ids\x18\x0b \x03(\x05\x12\x1c\n\x14zerofy_to_neuron_ids\x18\x0c \x03(\x05\x12+\n\x11zerofy_predicates\x18\r \x03(\x0e\x32\x10.ZerofyPredicateB\n\n\x08_op_typeB\x0b\n\t_layer_id\"_\n\x17WeightsOperationRequest\x12/\n\x10weight_operation\x18\x01 \x01(\x0b\x32\x10.WeightOperationH\x00\x88\x01\x01\x42\x13\n\x11_weight_operation\"<\n\x18WeightsOperationResponse\x12\x0f\n\x07success\x18\x01 \x01(\x08\x12\x0f\n\x07message\x18\x02 \x01(\t\"\xc1\x05\n\x0fHyperParameters\x12\x1c\n\x0f\x65xperiment_name\x18\x01 \x01(\tH\x00\x88\x01\x01\x12!\n\x14training_steps_to_do\x18\x02 \x01(\x05H\x01\x88\x01\x01\x12\x1a\n\rlearning_rate\x18\x03 \x01(\x02H\x02\x88\x01\x01\x12\x17\n\nbatch_size\x18\x04 \x01(\x05H\x03\x88\x01\x01\x12 \n\x13\x66ull_eval_frequency\x18\x05 \x01(\x05H\x04\x88\x01\x01\x12 \n\x13\x63heckpont_frequency\x18\x06 \x01(\x05H\x05\x88\x01\x01\x12\x18\n\x0bis_training\x18\x07 \x01(\x08H\x06\x88\x01\x01\x12\x15\n\x08nb_steps\x18\x08 \x01(\x05H\x07\x88\x01\x01\x12\x19\n\x0c\x61uditor_mode\x18\t \x01(\x08H\x08\x88\x01\x01\x12\x1d\n\x10train_batch_size\x18\n \x01(\x05H\t\x88\x01\x01\x12\x1b\n\x0eval_batch_size\x18\x0b \x01(\x05H\n\x88\x01\x01\x12\x1c\n\x0ftest_batch_size\x18\x0c \x01(\x05H\x0b\x88\x01\x01\x12\x1c\n\x0f\x65valuation_mode\x18\r \x01(\x08H\x0c\x88\x01\x01\x12\x1e\n\x11\x65valuation_config\x18\x0e \x01(\tH\r\x88\x01\x01\x42\x12\n\x10_experiment_nameB\x17\n\x15_training_steps_to_doB\x10\n\x0e_learning_rateB\r\n\x0b_batch_sizeB\x16\n\x14_full_eval_frequencyB\x16\n\x14_checkpont_frequencyB\x0e\n\x0c_is_trainingB\x0b\n\t_nb_stepsB\x0f\n\r_auditor_modeB\x13\n\x11_train_batch_sizeB\x11\n\x0f_val_batch_sizeB\x12\n\x10_test_batch_sizeB\x12\n\x10_evaluation_modeB\x14\n\x12_evaluation_config\",\n\rMetricsStatus\x12\x0c\n\x04name\x18\x01 \x01(\t\x12\r\n\x05value\x18\x02 \x01(\x02\"~\n\rAnnotatStatus\x12\x0c\n\x04name\x18\x01 \x01(\t\x12.\n\x08metadata\x18\x02 \x03(\x0b\x32\x1c.AnnotatStatus.MetadataEntry\x1a/\n\rMetadataEntry\x12\x0b\n\x03key\x18\x01 \x01(\t\x12\r\n\x05value\x18\x02 \x01(\x02:\x02\x38\x01\"\x90\x02\n\x10TrainingStatusEx\x12\x16\n\ttimestamp\x18\x01 \x01(\tH\x00\x88\x01\x01\x12\x1c\n\x0f\x65xperiment_name\x18\x02 \x01(\tH\x01\x88\x01\x01\x12\x16\n\tmodel_age\x18\x03 \x01(\x05H\x02\x88\x01\x01\x12+\n\x0emetrics_status\x18\x04 \x01(\x0b\x32\x0e.MetricsStatusH\x03\x88\x01\x01\x12+\n\x0e\x61nnotat_status\x18\x05 \x01(\x0b\x32\x0e.AnnotatStatusH\x04\x88\x01\x01\x42\x0c\n\n_timestampB\x12\n\x10_experiment_nameB\x0c\n\n_model_ageB\x11\n\x0f_metrics_statusB\x11\n\x0f_annotat_status\"]\n\x15HyperParameterCommand\x12/\n\x10hyper_parameters\x18\x01 \x01(\x0b\x32\x10.HyperParametersH\x00\x88\x01\x01\x42\x13\n\x11_hyper_parameters\">\n\x14\x44\x65nySamplesOperation\x12\x12\n\nsample_ids\x18\x01 \x03(\t\x12\x12\n\naccumulate\x18\x02 \x01(\x08\"0\n\x17LoadCheckpointOperation\x12\x15\n\rcheckpoint_id\x18\x01 \x01(\x05\"b\n\x11PlotNoteOperation\x12\x13\n\x0bmetric_name\x18\x01 \x01(\t\x12\x17\n\x0f\x65xperiment_hash\x18\x02 \x01(\t\x12\x11\n\tmodel_age\x18\x03 \x01(\x05\x12\x0c\n\x04note\x18\x04 \x01(\t\"L\n\x17SaveCheckpointOperation\x12\x19\n\x11save_architecture\x18\x01 \x01(\x08\x12\x16\n\x0esave_optimizer\x18\x02 \x01(\x08\"\x1a\n\x18RestartInstanceOperation\"\x8d\x08\n\x0eTrainerCommand\x12\x1c\n\x14get_hyper_parameters\x18\x04 \x01(\x08\x12\x1e\n\x16get_interactive_layers\x18\x05 \x01(\x08\x12\x1d\n\x10get_data_records\x18\x06 \x01(\tH\x00\x88\x01\x01\x12%\n\x18get_single_layer_info_id\x18\x08 \x01(\x05H\x01\x88\x01\x01\x12;\n\x16hyper_parameter_change\x18\x01 \x01(\x0b\x32\x16.HyperParameterCommandH\x02\x88\x01\x01\x12:\n\x16\x64\x65ny_samples_operation\x18\x07 \x01(\x0b\x32\x15.DenySamplesOperationH\x03\x88\x01\x01\x12?\n\x1b\x64\x65ny_eval_samples_operation\x18\n \x01(\x0b\x32\x15.DenySamplesOperationH\x04\x88\x01\x01\x12@\n\x19load_checkpoint_operation\x18\t \x01(\x0b\x32\x18.LoadCheckpointOperationH\x05\x88\x01\x01\x12\x42\n\x1eremove_from_denylist_operation\x18\x0b \x01(\x0b\x32\x15.DenySamplesOperationH\x06\x88\x01\x01\x12G\n#remove_eval_from_denylist_operation\x18\x0c \x01(\x0b\x32\x15.DenySamplesOperationH\x07\x88\x01\x01\x12\x34\n\x13plot_note_operation\x18\r \x01(\x0b\x32\x12.PlotNoteOperationH\x08\x88\x01\x01\x12@\n\x19save_checkpoint_operation\x18\x0e \x01(\x0b\x32\x18.SaveCheckpointOperationH\t\x88\x01\x01\x12\x39\n\x11restart_operation\x18\x0f \x01(\x0b\x32\x19.RestartInstanceOperationH\n\x88\x01\x01\x42\x13\n\x11_get_data_recordsB\x1b\n\x19_get_single_layer_info_idB\x19\n\x17_hyper_parameter_changeB\x19\n\x17_deny_samples_operationB\x1e\n\x1c_deny_eval_samples_operationB\x1c\n\x1a_load_checkpoint_operationB!\n\x1f_remove_from_denylist_operationB&\n$_remove_eval_from_denylist_operationB\x16\n\x14_plot_note_operationB\x1c\n\x1a_save_checkpoint_operationB\x14\n\x12_restart_operation\"\x9d\x01\n\x12HyperParameterDesc\x12\r\n\x05label\x18\x01 \x01(\t\x12\x0c\n\x04name\x18\x02 \x01(\t\x12\x0c\n\x04type\x18\x03 \x01(\t\x12\x1c\n\x0fnumerical_value\x18\x04 \x01(\x02H\x00\x88\x01\x01\x12\x19\n\x0cstring_value\x18\x05 \x01(\tH\x01\x88\x01\x01\x42\x12\n\x10_numerical_valueB\x0f\n\r_string_value\"\xf2\x02\n\x10NeuronStatistics\x12!\n\tneuron_id\x18\x01 \x01(\x0b\x32\t.NeuronIdH\x00\x88\x01\x01\x12\x17\n\nneuron_age\x18\x02 \x01(\x05H\x01\x88\x01\x01\x12\x1f\n\x12train_trigger_rate\x18\x03 \x01(\x02H\x02\x88\x01\x01\x12\x1e\n\x11\x65val_trigger_rate\x18\x04 \x01(\x02H\x03\x88\x01\x01\x12\x1a\n\rlearning_rate\x18\x07 \x01(\x02H\x04\x88\x01\x01\x12\x36\n\x0bincoming_lr\x18\x08 \x03(\x0b\x32!.NeuronStatistics.IncomingLrEntry\x1a\x31\n\x0fIncomingLrEntry\x12\x0b\n\x03key\x18\x01 \x01(\x05\x12\r\n\x05value\x18\x02 \x01(\x02:\x02\x38\x01\x42\x0c\n\n_neuron_idB\r\n\x0b_neuron_ageB\x15\n\x13_train_trigger_rateB\x14\n\x12_eval_trigger_rateB\x10\n\x0e_learning_rate\"\xf0\x02\n\x13LayerRepresentation\x12\x15\n\x08layer_id\x18\x01 \x01(\x05H\x00\x88\x01\x01\x12\x17\n\nlayer_name\x18\x02 \x01(\tH\x01\x88\x01\x01\x12\x17\n\nlayer_type\x18\x03 \x01(\tH\x02\x88\x01\x01\x12\x1a\n\rneurons_count\x18\x04 \x01(\x05H\x03\x88\x01\x01\x12#\n\x16incoming_neurons_count\x18\x05 \x01(\x05H\x04\x88\x01\x01\x12\x18\n\x0bkernel_size\x18\x06 \x01(\x05H\x05\x88\x01\x01\x12\x13\n\x06stride\x18\x07 \x01(\x05H\x06\x88\x01\x01\x12-\n\x12neurons_statistics\x18\n \x03(\x0b\x32\x11.NeuronStatisticsB\x0b\n\t_layer_idB\r\n\x0b_layer_nameB\r\n\x0b_layer_typeB\x10\n\x0e_neurons_countB\x19\n\x17_incoming_neurons_countB\x0e\n\x0c_kernel_sizeB\t\n\x07_stride\"H\n\x11\x41\x63tivationRequest\x12\x10\n\x08layer_id\x18\x01 \x01(\x05\x12\x11\n\tsample_id\x18\x02 \x01(\t\x12\x0e\n\x06origin\x18\x03 \x01(\t\"H\n\rActivationMap\x12\x11\n\tneuron_id\x18\x01 \x01(\x05\x12\x0e\n\x06values\x18\x02 \x03(\x02\x12\t\n\x01H\x18\x03 \x01(\x05\x12\t\n\x01W\x18\x04 \x01(\x05\"d\n\x12\x41\x63tivationResponse\x12\x12\n\nlayer_type\x18\x01 \x01(\t\x12\x15\n\rneurons_count\x18\x02 \x01(\x05\x12#\n\x0b\x61\x63tivations\x18\x03 \x03(\x0b\x32\x0e.ActivationMap\"\x93\x01\n\tTaskField\x12\x0c\n\x04name\x18\x01 \x01(\t\x12\x15\n\x0b\x66loat_value\x18\x02 \x01(\x02H\x00\x12\x13\n\tint_value\x18\x03 \x01(\x05H\x00\x12\x16\n\x0cstring_value\x18\x04 \x01(\tH\x00\x12\x15\n\x0b\x62ytes_value\x18\x05 \x01(\x0cH\x00\x12\x14\n\nbool_value\x18\x06 \x01(\x08H\x00\x42\x07\n\x05value\"\x87\x03\n\x0eRecordMetadata\x12\x11\n\tsample_id\x18\x01 \x01(\t\x12\x14\n\x0csample_label\x18\x02 \x03(\x05\x12\x19\n\x11sample_prediction\x18\x03 \x03(\x05\x12=\n\x10sample_last_loss\x18\x04 \x03(\x0b\x32#.RecordMetadata.SampleLastLossEntry\x12\x19\n\x11sample_encounters\x18\x05 \x01(\x05\x12\x18\n\x10sample_discarded\x18\x06 \x01(\x08\x12 \n\x0c\x65xtra_fields\x18\x07 \x03(\x0b\x32\n.TaskField\x12\x16\n\x0eprediction_raw\x18\t \x01(\x0c\x12\x11\n\ttask_type\x18\n \x01(\t\x12\x19\n\x11sample_label_text\x18\x0b \x03(\t\x12\x1e\n\x16sample_prediction_text\x18\x0c \x03(\t\x1a\x35\n\x13SampleLastLossEntry\x12\x0b\n\x03key\x18\x01 \x01(\t\x12\r\n\x05value\x18\x02 \x01(\x02:\x02\x38\x01\"\x93\x01\n\x10SampleStatistics\x12\x13\n\x06origin\x18\x06 \x01(\tH\x00\x88\x01\x01\x12\x19\n\x0csample_count\x18\x07 \x01(\x05H\x01\x88\x01\x01\x12\x11\n\ttask_type\x18\t \x01(\t\x12 \n\x07records\x18\x08 \x03(\x0b\x32\x0f.RecordMetadataB\t\n\x07_originB\x0f\n\r_sample_count\"\xe6\x01\n\x0f\x43ommandResponse\x12\x0f\n\x07success\x18\x01 \x01(\x08\x12\x0f\n\x07message\x18\x02 \x01(\t\x12\x33\n\x16hyper_parameters_descs\x18\x03 \x03(\x0b\x32\x13.HyperParameterDesc\x12\x33\n\x15layer_representations\x18\x04 \x03(\x0b\x32\x14.LayerRepresentation\x12\x31\n\x11sample_statistics\x18\x05 \x01(\x0b\x32\x11.SampleStatisticsH\x00\x88\x01\x01\x42\x14\n\x12_sample_statistics\"U\n\rSampleRequest\x12\x16\n\tsample_id\x18\x01 \x01(\tH\x00\x88\x01\x01\x12\x13\n\x06origin\x18\x02 \x01(\tH\x01\x88\x01\x01\x42\x0c\n\n_sample_idB\t\n\x07_origin\"\xad\x02\n\x15SampleRequestResponse\x12\x16\n\tsample_id\x18\x01 \x01(\tH\x00\x88\x01\x01\x12\x13\n\x06origin\x18\x02 \x01(\tH\x01\x88\x01\x01\x12\x12\n\x05label\x18\x03 \x01(\x05H\x02\x88\x01\x01\x12\x11\n\x04\x64\x61ta\x18\x04 \x01(\x0cH\x03\x88\x01\x01\x12\x1a\n\rerror_message\x18\x05 \x01(\tH\x04\x88\x01\x01\x12\x15\n\x08raw_data\x18\x06 \x01(\x0cH\x05\x88\x01\x01\x12\x11\n\x04mask\x18\x07 \x01(\x0cH\x06\x88\x01\x01\x12\x17\n\nprediction\x18\x08 \x01(\x0cH\x07\x88\x01\x01\x42\x0c\n\n_sample_idB\t\n\x07_originB\x08\n\x06_labelB\x07\n\x05_dataB\x10\n\x0e_error_messageB\x0b\n\t_raw_dataB\x07\n\x05_maskB\r\n\x0b_prediction\"\x92\x01\n\x12\x42\x61tchSampleRequest\x12\x12\n\nsample_ids\x18\x01 \x03(\t\x12\x0e\n\x06origin\x18\x02 \x01(\t\x12\x19\n\x0cresize_width\x18\x03 \x01(\x05H\x00\x88\x01\x01\x12\x1a\n\rresize_height\x18\x04 \x01(\x05H\x01\x88\x01\x01\x42\x0f\n\r_resize_widthB\x10\n\x0e_resize_height\">\n\x13\x42\x61tchSampleResponse\x12\'\n\x07samples\x18\x01 \x03(\x0b\x32\x16.SampleRequestResponse\".\n\x0eWeightsRequest\x12\x1c\n\tneuron_id\x18\x01 \x01(\x0b\x32\t.NeuronId\"\x9d\x02\n\x0fWeightsResponse\x12\x1c\n\tneuron_id\x18\x01 \x01(\x0b\x32\t.NeuronId\x12\x17\n\nlayer_name\x18\x02 \x01(\tH\x00\x88\x01\x01\x12\x17\n\nlayer_type\x18\x03 \x01(\tH\x01\x88\x01\x01\x12\x10\n\x08incoming\x18\x04 \x01(\x05\x12\x10\n\x08outgoing\x18\x05 \x01(\x05\x12\x18\n\x0bkernel_size\x18\x06 \x01(\x05H\x02\x88\x01\x01\x12\x0f\n\x07weights\x18\x07 \x03(\x02\x12\x0f\n\x07success\x18\x0b \x01(\x08\x12\x1a\n\rerror_message\x18\x0c \x01(\tH\x03\x88\x01\x01\x42\r\n\x0b_layer_nameB\r\n\x0b_layer_typeB\x0e\n\x0c_kernel_sizeB\x10\n\x0e_error_message\"R\n\x10\x44\x61taQueryRequest\x12\r\n\x05query\x18\x01 \x01(\t\x12\x12\n\naccumulate\x18\x02 \x01(\x08\x12\x1b\n\x13is_natural_language\x18\x03 \x01(\x08\"5\n\x11\x43\x61tegoricalTagDef\x12\x0c\n\x04name\x18\x01 \x01(\t\x12\x12\n\ncategories\x18\x02 \x03(\t\"\xa9\x02\n\x11\x44\x61taQueryResponse\x12\x0f\n\x07success\x18\x01 \x01(\x08\x12\x0f\n\x07message\x18\x02 \x01(\t\x12\x1d\n\x15number_of_all_samples\x18\x03 \x01(\x05\x12%\n\x1dnumber_of_samples_in_the_loop\x18\x04 \x01(\x05\x12#\n\x1bnumber_of_discarded_samples\x18\x05 \x01(\x05\x12\x13\n\x0bunique_tags\x18\x06 \x03(\t\x12+\n\x11\x61gent_intent_type\x18\x07 \x01(\x0e\x32\x10.AgentIntentType\x12\x17\n\x0f\x61nalysis_result\x18\x08 \x01(\t\x12,\n\x10\x63\x61tegorical_tags\x18\t \x03(\x0b\x32\x12.CategoricalTagDef\"\xc2\x01\n\x12\x44\x61taSamplesRequest\x12\x13\n\x0bstart_index\x18\x01 \x01(\x05\x12\x13\n\x0brecords_cnt\x18\x02 \x01(\x05\x12 \n\x18include_transformed_data\x18\x03 \x01(\x08\x12\x18\n\x10include_raw_data\x18\x04 \x01(\x08\x12\x19\n\x11stats_to_retrieve\x18\x05 \x03(\t\x12\x14\n\x0cresize_width\x18\x06 \x01(\x05\x12\x15\n\rresize_height\x18\x07 \x01(\x05\"m\n\x08\x44\x61taStat\x12\x0c\n\x04name\x18\x01 \x01(\t\x12\x0c\n\x04type\x18\x02 \x01(\t\x12\r\n\x05shape\x18\x03 \x03(\x05\x12\r\n\x05value\x18\x04 \x03(\x02\x12\x14\n\x0cvalue_string\x18\x05 \x01(\t\x12\x11\n\tthumbnail\x18\x06 \x01(\x0c\">\n\nDataRecord\x12\x11\n\tsample_id\x18\x01 \x01(\t\x12\x1d\n\ndata_stats\x18\x02 \x03(\x0b\x32\t.DataStat\"\x97\x01\n\x13\x44\x61taSamplesResponse\x12\x0f\n\x07success\x18\x01 \x01(\x08\x12\x0f\n\x07message\x18\x02 \x01(\t\x12!\n\x0c\x64\x61ta_records\x18\x03 \x03(\x0b\x32\x0b.DataRecord\x12\x12\n\nis_subview\x18\x04 \x01(\x08\x12\x12\n\nview_count\x18\x05 \x01(\x03\x12\x13\n\x0btotal_count\x18\x06 \x01(\x03\"C\n\x0fHistogramSubBar\x12\x0e\n\x06origin\x18\x01 \x01(\t\x12\x11\n\tdiscarded\x18\x02 \x01(\x08\x12\r\n\x05\x63ount\x18\x03 \x01(\x03\"h\n\x0cHistogramBin\x12\x0b\n\x03min\x18\x01 \x01(\x01\x12\x0b\n\x03max\x18\x02 \x01(\x01\x12\x0b\n\x03\x61vg\x18\x03 \x01(\x01\x12\r\n\x05\x63ount\x18\x04 \x01(\x03\x12\"\n\x08sub_bars\x18\x05 \x03(\x0b\x32\x10.HistogramSubBar\"[\n\x17\x43\x61tegoricalHistogramBar\x12\r\n\x05label\x18\x01 \x01(\t\x12\r\n\x05\x63ount\x18\x02 \x01(\x03\x12\"\n\x08sub_bars\x18\x03 \x03(\x0b\x32\x10.HistogramSubBar\"4\n\x10HistogramRequest\x12\x0e\n\x06\x63olumn\x18\x01 \x01(\t\x12\x10\n\x08max_bins\x18\x02 \x01(\x05\"\xb2\x01\n\x11HistogramResponse\x12\x0f\n\x07success\x18\x01 \x01(\x08\x12\x0f\n\x07message\x18\x02 \x01(\t\x12\x12\n\ntotal_rows\x18\x03 \x01(\x03\x12\x1b\n\x04\x62ins\x18\x04 \x03(\x0b\x32\r.HistogramBin\x12\x16\n\x0eis_categorical\x18\x05 \x01(\x08\x12\x32\n\x10\x63\x61tegorical_bars\x18\x06 \x03(\x0b\x32\x18.CategoricalHistogramBar\"W\n\x12GetMetaDataRequest\x12\x13\n\x0bstart_index\x18\x01 \x01(\x05\x12\x13\n\x0brecords_cnt\x18\x02 \x01(\x05\x12\x17\n\x0fmodal_sample_id\x18\x03 \x01(\t\"\x99\x01\n\x13GetMetaDataResponse\x12\x0f\n\x07success\x18\x01 \x01(\x08\x12\x0f\n\x07message\x18\x02 \x01(\t\x12\x1a\n\x12\x61ll_metadata_names\x18\x03 \x03(\t\x12!\n\x0cgrid_records\x18\x04 \x03(\x0b\x32\x0b.DataRecord\x12!\n\x0cmodal_record\x18\x05 \x01(\x0b\x32\x0b.DataRecord\"j\n\x12StepSamplesRequest\x12\x13\n\x0bmetric_name\x18\x01 \x01(\t\x12\x17\n\x0f\x65xperiment_hash\x18\x02 \x01(\t\x12\x11\n\tmodel_age\x18\x03 \x01(\x05\x12\x13\n\x0bmax_samples\x18\x04 \x01(\x05\"{\n\x13StepSamplesResponse\x12\x0f\n\x07success\x18\x01 \x01(\x08\x12\x0f\n\x07message\x18\x02 \x01(\t\x12\x12\n\nsample_ids\x18\x03 \x03(\t\x12\x17\n\x0ftotal_available\x18\x04 \x01(\x05\x12\x15\n\rsample_values\x18\x05 \x03(\x02\"Y\n\x1aGetSignalTrajectoryRequest\x12\x13\n\x0bsignal_name\x18\x01 \x01(\t\x12\x12\n\nsample_ids\x18\x02 \x03(\t\x12\x12\n\nmax_points\x18\x03 \x01(\x05\"4\n\x10SignalTrajectory\x12\x11\n\tsample_id\x18\x01 \x01(\t\x12\r\n\x05value\x18\x02 \x03(\x02\"}\n\x1bGetSignalTrajectoryResponse\x12\x0f\n\x07success\x18\x01 \x01(\x08\x12\x0f\n\x07message\x18\x02 \x01(\t\x12\x13\n\x0bsignal_name\x18\x03 \x01(\t\x12\'\n\x0ctrajectories\x18\x04 \x03(\x0b\x32\x11.SignalTrajectory\"Y\n\x11PointCloudRequest\x12\x11\n\tsample_id\x18\x01 \x01(\t\x12\x0e\n\x06origin\x18\x02 \x01(\t\x12\x12\n\nmax_points\x18\x03 \x01(\x05\x12\r\n\x05\x66ield\x18\x04 \x01(\t\"\xbf\x01\n\x0fPointCloudChunk\x12\x0f\n\x07success\x18\x01 \x01(\x08\x12\x0f\n\x07message\x18\x02 \x01(\t\x12\x12\n\nnum_points\x18\x03 \x01(\x05\x12\x14\n\x0cnum_features\x18\x04 \x01(\x05\x12\x10\n\x08pc_range\x18\x05 \x03(\x02\x12\x0c\n\x04\x64\x61ta\x18\x06 \x01(\x0c\x12\x13\n\x0b\x63hunk_index\x18\x07 \x01(\x05\x12\x14\n\x0ctotal_chunks\x18\x08 \x01(\x05\x12\x15\n\rfeature_names\x18\t \x03(\t\"b\n\x0cMediaRequest\x12\x11\n\tsample_id\x18\x01 \x01(\t\x12\x0e\n\x06origin\x18\x02 \x01(\t\x12\x0c\n\x04kind\x18\x03 \x01(\t\x12\x12\n\nmax_frames\x18\x04 \x01(\x05\x12\r\n\x05\x66ield\x18\x05 \x01(\t\"\x92\x02\n\nMediaChunk\x12\x0f\n\x07success\x18\x01 \x01(\x08\x12\x0f\n\x07message\x18\x02 \x01(\t\x12\x11\n\tmime_type\x18\x03 \x01(\t\x12\x13\n\x0b\x66rame_count\x18\x04 \x01(\x05\x12\x0b\n\x03\x66ps\x18\x05 \x01(\x02\x12\x11\n\thas_audio\x18\x06 \x01(\x08\x12\r\n\x05width\x18\x07 \x01(\x05\x12\x0e\n\x06height\x18\x08 \x01(\x05\x12\x13\n\x0btotal_bytes\x18\t \x01(\x05\x12\x18\n\x10\x64uration_seconds\x18\n \x01(\x02\x12\x13\n\x0bsample_rate\x18\x0b \x01(\x05\x12\x0c\n\x04\x64\x61ta\x18\x0c \x01(\x0c\x12\x13\n\x0b\x63hunk_index\x18\r \x01(\x05\x12\x14\n\x0ctotal_chunks\x18\x0e \x01(\x05\"\x8c\x02\n\x10\x44\x61taEditsRequest\x12\x11\n\tstat_name\x18\x01 \x01(\t\x12\x13\n\x0b\x66loat_value\x18\x02 \x01(\x02\x12\x14\n\x0cstring_value\x18\x03 \x01(\t\x12\x12\n\nbool_value\x18\x04 \x01(\x08\x12\x1d\n\x04type\x18\x05 \x01(\x0e\x32\x0f.SampleEditType\x12\x13\n\x0bsamples_ids\x18\x06 \x03(\t\x12\x16\n\x0esample_origins\x18\x07 \x03(\t\x12\x16\n\x0eis_categorical\x18\x08 \x01(\x08\x12\x12\n\ncategories\x18\t \x03(\t\x12\x15\n\rsample_values\x18\n \x03(\x02\x12\x17\n\x0f\x65xperiment_hash\x18\x0b \x01(\t\"5\n\x11\x44\x61taEditsResponse\x12\x0f\n\x07success\x18\x01 \x01(\x08\x12\x0f\n\x07message\x18\x02 \x01(\t\":\n\x12\x44\x61taSplitsResponse\x12\x0f\n\x07success\x18\x01 \x01(\x08\x12\x13\n\x0bsplit_names\x18\x02 \x03(\t\"9\n\x13\x41gentHealthResponse\x12\x11\n\tavailable\x18\x01 \x01(\x08\x12\x0f\n\x07message\x18\x02 \x01(\t\"^\n\x16InitializeAgentRequest\x12\x0f\n\x07\x61pi_key\x18\x01 \x01(\t\x12$\n\x08provider\x18\x02 \x01(\x0e\x32\x12.AgentProviderType\x12\r\n\x05model\x18\x03 \x01(\t\";\n\x17InitializeAgentResponse\x12\x0f\n\x07success\x18\x01 \x01(\x08\x12\x0f\n\x07message\x18\x02 \x01(\t\"(\n\x17\x43hangeAgentModelRequest\x12\r\n\x05model\x18\x01 \x01(\t\"<\n\x18\x43hangeAgentModelResponse\x12\x0f\n\x07success\x18\x01 \x01(\x08\x12\x0f\n\x07message\x18\x02 \x01(\t\"\x17\n\x15GetAgentModelsRequest\"J\n\x16GetAgentModelsResponse\x12\x0f\n\x07success\x18\x01 \x01(\x08\x12\x0e\n\x06models\x18\x02 \x03(\t\x12\x0f\n\x07message\x18\x03 \x01(\t\"6\n\x12ResetAgentResponse\x12\x0f\n\x07success\x18\x01 \x01(\x08\x12\x0f\n\x07message\x18\x02 \x01(\t\"=\n\x19\x43learAgentHistoryResponse\x12\x0f\n\x07success\x18\x01 \x01(\x08\x12\x0f\n\x07message\x18\x02 \x01(\t\"?\n\x1b\x43ompactAgentHistoryResponse\x12\x0f\n\x07success\x18\x01 \x01(\x08\x12\x0f\n\x07message\x18\x02 \x01(\t\"\xe5\x01\n\x1cGetAgentContextUsageResponse\x12\x0f\n\x07success\x18\x01 \x01(\x08\x12\x0f\n\x07message\x18\x02 \x01(\t\x12\r\n\x05model\x18\x03 \x01(\t\x12\x16\n\x0e\x63ontext_window\x18\x04 \x01(\x03\x12\x14\n\x0cinput_tokens\x18\x05 \x01(\x03\x12\x15\n\routput_tokens\x18\x06 \x01(\x03\x12\x18\n\x10reasoning_tokens\x18\x07 \x01(\x03\x12\x19\n\x11\x63\x61\x63he_read_tokens\x18\x08 \x01(\x03\x12\x1a\n\x12\x63\x61\x63he_write_tokens\x18\t \x01(\x03\"3\n\x18RestoreCheckpointRequest\x12\x17\n\x0f\x65xperiment_hash\x18\x01 \x01(\t\"=\n\x19RestoreCheckpointResponse\x12\x0f\n\x07success\x18\x01 \x01(\x08\x12\x0f\n\x07message\x18\x02 \x01(\t\"\xa8\x01\n\x11\x45xperimentRunInfo\x12\x17\n\x0f\x65xperiment_hash\x18\x01 \x01(\t\x12\x17\n\x0f\x65xperiment_name\x18\x02 \x01(\t\x12\r\n\x05notes\x18\x03 \x01(\t\x12\x0f\n\x07\x63reated\x18\x04 \x01(\t\x12\x11\n\tlast_used\x18\x05 \x01(\t\x12\x1a\n\x12latest_weight_step\x18\x06 \x01(\x05\x12\x12\n\nis_current\x18\x07 \x01(\x08\"\x1b\n\x19ListExperimentRunsRequest\">\n\x1aListExperimentRunsResponse\x12 \n\x04runs\x18\x01 \x03(\x0b\x32\x12.ExperimentRunInfo\"G\n\x1aRenameExperimentRunRequest\x12\x17\n\x0f\x65xperiment_hash\x18\x01 \x01(\t\x12\x10\n\x08new_name\x18\x02 \x01(\t\"?\n\x1bRenameExperimentRunResponse\x12\x0f\n\x07success\x18\x01 \x01(\x08\x12\x0f\n\x07message\x18\x02 \x01(\t\"F\n\x1cSetExperimentRunNotesRequest\x12\x17\n\x0f\x65xperiment_hash\x18\x01 \x01(\t\x12\r\n\x05notes\x18\x02 \x01(\t\"A\n\x1dSetExperimentRunNotesResponse\x12\x0f\n\x07success\x18\x01 \x01(\x08\x12\x0f\n\x07message\x18\x02 \x01(\t\"R\n\x18TriggerEvaluationRequest\x12\x12\n\nsplit_name\x18\x01 \x01(\t\x12\x0c\n\x04tags\x18\x02 \x03(\t\x12\x14\n\x0cuse_full_set\x18\x03 \x01(\x08\"=\n\x19TriggerEvaluationResponse\x12\x0f\n\x07success\x18\x01 \x01(\x08\x12\x0f\n\x07message\x18\x02 \x01(\t\"\x1c\n\x1aGetEvaluationStatusRequest\"\x81\x01\n\x1bGetEvaluationStatusResponse\x12\x0e\n\x06status\x18\x01 \x01(\t\x12\x0f\n\x07\x63urrent\x18\x02 \x01(\x05\x12\r\n\x05total\x18\x03 \x01(\x05\x12\x0f\n\x07message\x18\x04 \x01(\t\x12\r\n\x05\x65rror\x18\x05 \x01(\t\x12\x12\n\nsplit_name\x18\x06 \x01(\t\")\n\x17\x43\x61ncelEvaluationRequest\x12\x0e\n\x06reason\x18\x01 \x01(\t\"<\n\x18\x43\x61ncelEvaluationResponse\x12\x0f\n\x07success\x18\x01 \x01(\x08\x12\x0f\n\x07message\x18\x02 \x01(\t\"7\n\x16RunNotebookCellRequest\x12\x0c\n\x04\x63ode\x18\x01 \x01(\t\x12\x0f\n\x07\x63\x65ll_id\x18\x02 \x01(\t\"2\n\x10NotebookCellDone\x12\x12\n\nexec_count\x18\x01 \x01(\x05\x12\n\n\x02ok\x18\x02 \x01(\x08\"\x1e\n\x1cInterruptNotebookCellRequest\":\n\x1dInterruptNotebookCellResponse\x12\n\n\x02ok\x18\x01 \x01(\x08\x12\r\n\x05\x65rror\x18\x02 \x01(\t\"\xbd\x01\n\x11NotebookCellChunk\x12\x0f\n\x07\x63\x65ll_id\x18\x01 \x01(\t\x12\x10\n\x06stdout\x18\x02 \x01(\tH\x00\x12\x10\n\x06stderr\x18\x03 \x01(\tH\x00\x12\x15\n\x0bresult_text\x18\x04 \x01(\tH\x00\x12\x13\n\timage_png\x18\x05 \x01(\x0cH\x00\x12\x19\n\x0f\x65rror_traceback\x18\x06 \x01(\tH\x00\x12!\n\x04\x64one\x18\x07 \x01(\x0b\x32\x11.NotebookCellDoneH\x00\x42\t\n\x07payload\"S\n\x10NotebookResponse\x12\x12\n\nipynb_json\x18\x01 \x01(\t\x12\x0f\n\x07\x65xisted\x18\x02 \x01(\x08\x12\x0c\n\x04path\x18\x03 \x01(\t\x12\x0c\n\x04name\x18\x04 \x01(\t\"7\n\x13SaveNotebookRequest\x12\x12\n\nipynb_json\x18\x01 \x01(\t\x12\x0c\n\x04name\x18\x02 \x01(\t\"M\n\x14SaveNotebookResponse\x12\n\n\x02ok\x18\x01 \x01(\x08\x12\x0c\n\x04path\x18\x02 \x01(\t\x12\r\n\x05\x65rror\x18\x03 \x01(\t\x12\x0c\n\x04name\x18\x04 \x01(\t\"C\n\x1bGenerateNotebookCodeRequest\x12\x0e\n\x06prompt\x18\x01 \x01(\t\x12\x14\n\x0c\x63ontext_code\x18\x02 \x01(\t\"\\\n\x1cGenerateNotebookCodeResponse\x12\x0c\n\x04\x63ode\x18\x01 \x01(\t\x12\x13\n\x0b\x65xplanation\x18\x02 \x01(\t\x12\n\n\x02ok\x18\x03 \x01(\x08\x12\r\n\x05\x65rror\x18\x04 \x01(\t\"~\n\x18\x45xportAnnotationsRequest\x12\'\n\x06\x66ormat\x18\x01 \x01(\x0e\x32\x17.AnnotationExportFormat\x12\x0e\n\x06origin\x18\x02 \x01(\t\x12\x1b\n\x13include_predictions\x18\x03 \x01(\x08\x12\x0c\n\x04tags\x18\x04 \x03(\t\"\x88\x01\n\x19\x45xportAnnotationsResponse\x12\x0f\n\x07success\x18\x01 \x01(\x08\x12\x0f\n\x07message\x18\x02 \x01(\t\x12\x0f\n\x07payload\x18\x03 \x01(\x0c\x12\x10\n\x08\x66ilename\x18\x04 \x01(\t\x12\x11\n\tmime_type\x18\x05 \x01(\t\x12\x13\n\x0bimage_count\x18\x06 \x01(\x05*d\n\x13WeightOperationType\x12\n\n\x06ZEROFY\x10\x00\x12\x10\n\x0cREINITIALIZE\x10\x01\x12\n\n\x06\x46REEZE\x10\x02\x12\x12\n\x0eREMOVE_NEURONS\x10\t\x12\x0f\n\x0b\x41\x44\x44_NEURONS\x10\n*o\n\x0fZerofyPredicate\x12\x19\n\x15ZEROFY_PREDICATE_NONE\x10\x00\x12 \n\x1cZEROFY_PREDICATE_WITH_FROZEN\x10\x01\x12\x1f\n\x1bZEROFY_PREDICATE_WITH_OLDER\x10\x02*M\n\x0f\x41gentIntentType\x12\x12\n\x0eINTENT_UNKNOWN\x10\x00\x12\x11\n\rINTENT_FILTER\x10\x01\x12\x13\n\x0fINTENT_ANALYSIS\x10\x02*I\n\x0eSampleEditType\x12\x11\n\rEDIT_OVERRIDE\x10\x00\x12\x13\n\x0f\x45\x44IT_ACCUMULATE\x10\x01\x12\x0f\n\x0b\x45\x44IT_REMOVE\x10\x02*C\n\x11\x41gentProviderType\x12\x17\n\x13PROVIDER_OPENROUTER\x10\x00\x12\x15\n\x11PROVIDER_OPENCODE\x10\x01*m\n\x16\x41nnotationExportFormat\x12\x16\n\x12\x45XPORT_FORMAT_CVAT\x10\x00\x12\x1e\n\x1a\x45XPORT_FORMAT_LABEL_STUDIO\x10\x01\x12\x1b\n\x17\x45XPORT_FORMAT_V7_DARWIN\x10\x02\x32\xfe\x12\n\x11\x45xperimentService\x12P\n\x13GetLatestLoggerData\x12\x1b.GetLatestLoggerDataRequest\x1a\x1c.GetLatestLoggerDataResponse\x12\x36\n\x11\x45xperimentCommand\x12\x0f.TrainerCommand\x1a\x10.CommandResponse\x12H\n\x11ManipulateWeights\x12\x18.WeightsOperationRequest\x1a\x19.WeightsOperationResponse\x12/\n\nGetWeights\x12\x0f.WeightsRequest\x1a\x10.WeightsResponse\x12\x39\n\x0eGetActivations\x12\x12.ActivationRequest\x1a\x13.ActivationResponse\x12\x37\n\nGetSamples\x12\x13.BatchSampleRequest\x1a\x14.BatchSampleResponse\x12\x37\n\x0e\x41pplyDataQuery\x12\x11.DataQueryRequest\x1a\x12.DataQueryResponse\x12;\n\x0eGetDataSamples\x12\x13.DataSamplesRequest\x1a\x14.DataSamplesResponse\x12\x35\n\x0cGetHistogram\x12\x11.HistogramRequest\x1a\x12.HistogramResponse\x12\x38\n\x0bGetMetaData\x12\x13.GetMetaDataRequest\x1a\x14.GetMetaDataResponse\x12P\n\x13GetSignalTrajectory\x12\x1b.GetSignalTrajectoryRequest\x1a\x1c.GetSignalTrajectoryResponse\x12;\n\x0eGetStepSamples\x12\x13.StepSamplesRequest\x1a\x14.StepSamplesResponse\x12\x37\n\rGetPointCloud\x12\x12.PointCloudRequest\x1a\x10.PointCloudChunk0\x01\x12(\n\x08GetMedia\x12\r.MediaRequest\x1a\x0b.MediaChunk0\x01\x12\x37\n\x0e\x45\x64itDataSample\x12\x11.DataEditsRequest\x1a\x12.DataEditsResponse\x12,\n\rGetDataSplits\x12\x06.Empty\x1a\x13.DataSplitsResponse\x12\x30\n\x10\x43heckAgentHealth\x12\x06.Empty\x1a\x14.AgentHealthResponse\x12\x44\n\x0fInitializeAgent\x12\x17.InitializeAgentRequest\x1a\x18.InitializeAgentResponse\x12G\n\x10\x43hangeAgentModel\x12\x18.ChangeAgentModelRequest\x1a\x19.ChangeAgentModelResponse\x12\x41\n\x0eGetAgentModels\x12\x16.GetAgentModelsRequest\x1a\x17.GetAgentModelsResponse\x12)\n\nResetAgent\x12\x06.Empty\x1a\x13.ResetAgentResponse\x12\x37\n\x11\x43learAgentHistory\x12\x06.Empty\x1a\x1a.ClearAgentHistoryResponse\x12;\n\x13\x43ompactAgentHistory\x12\x06.Empty\x1a\x1c.CompactAgentHistoryResponse\x12=\n\x14GetAgentContextUsage\x12\x06.Empty\x1a\x1d.GetAgentContextUsageResponse\x12@\n\x0fRunNotebookCell\x12\x17.RunNotebookCellRequest\x1a\x12.NotebookCellChunk0\x01\x12V\n\x15InterruptNotebookCell\x12\x1d.InterruptNotebookCellRequest\x1a\x1e.InterruptNotebookCellResponse\x12(\n\x0bGetNotebook\x12\x06.Empty\x1a\x11.NotebookResponse\x12;\n\x0cSaveNotebook\x12\x14.SaveNotebookRequest\x1a\x15.SaveNotebookResponse\x12S\n\x14GenerateNotebookCode\x12\x1c.GenerateNotebookCodeRequest\x1a\x1d.GenerateNotebookCodeResponse\x12J\n\x11RestoreCheckpoint\x12\x19.RestoreCheckpointRequest\x1a\x1a.RestoreCheckpointResponse\x12M\n\x12ListExperimentRuns\x12\x1a.ListExperimentRunsRequest\x1a\x1b.ListExperimentRunsResponse\x12P\n\x13RenameExperimentRun\x12\x1b.RenameExperimentRunRequest\x1a\x1c.RenameExperimentRunResponse\x12V\n\x15SetExperimentRunNotes\x12\x1d.SetExperimentRunNotesRequest\x1a\x1e.SetExperimentRunNotesResponse\x12J\n\x11TriggerEvaluation\x12\x19.TriggerEvaluationRequest\x1a\x1a.TriggerEvaluationResponse\x12P\n\x13GetEvaluationStatus\x12\x1b.GetEvaluationStatusRequest\x1a\x1c.GetEvaluationStatusResponse\x12G\n\x10\x43\x61ncelEvaluation\x12\x18.CancelEvaluationRequest\x1a\x19.CancelEvaluationResponse\x12J\n\x11\x45xportAnnotations\x12\x19.ExportAnnotationsRequest\x1a\x1a.ExportAnnotationsResponseb\x06proto3') +DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile(b'\n)weightslab/proto/experiment_service.proto\"\x81\x02\n\x1aGetLatestLoggerDataRequest\x12\x1c\n\x14request_full_history\x18\x01 \x01(\x08\x12\x12\n\nmax_points\x18\x02 \x01(\x05\x12\x17\n\x0f\x62reak_by_slices\x18\x03 \x01(\x08\x12\x0c\n\x04tags\x18\x04 \x03(\t\x12\x12\n\ngraph_name\x18\x05 \x01(\t\x12\r\n\x05x_min\x18\x06 \x01(\x03\x12\r\n\x05x_max\x18\x07 \x01(\x03\x12\x13\n\x0bhas_x_range\x18\x08 \x01(\x08\x12\x14\n\x0cmetric_names\x18\t \x03(\t\x12\x19\n\x11\x65xperiment_hashes\x18\n \x03(\t\x12\x12\n\nindex_only\x18\x0b \x01(\x08\"\xa2\x01\n\x10SignalCurveIndex\x12\x13\n\x0bmetric_name\x18\x01 \x01(\t\x12\x17\n\x0f\x65xperiment_hash\x18\x02 \x01(\t\x12\x12\n\nfirst_step\x18\x03 \x01(\x03\x12\x11\n\tlast_step\x18\x04 \x01(\x03\x12\x13\n\x0bpoint_count\x18\x05 \x01(\x03\x12\x11\n\tvalue_min\x18\x06 \x01(\x01\x12\x11\n\tvalue_max\x18\x07 \x01(\x01\"1\n\rSignalOutlier\x12\x11\n\tsample_id\x18\x01 \x01(\t\x12\r\n\x05value\x18\x02 \x01(\x02\"\xd2\x03\n\x0fLoggerDataPoint\x12\x13\n\x0bmetric_name\x18\x01 \x01(\t\x12\x11\n\tmodel_age\x18\x02 \x01(\x05\x12\x14\n\x0cmetric_value\x18\x03 \x01(\x02\x12\x17\n\x0f\x65xperiment_hash\x18\x04 \x01(\t\x12\x11\n\ttimestamp\x18\x05 \x01(\x03\x12\x11\n\tsample_id\x18\x06 \x01(\t\x12\x1c\n\x14is_evaluation_marker\x18\x07 \x01(\x08\x12\x12\n\nsplit_name\x18\x08 \x01(\t\x12\x17\n\x0f\x65valuation_tags\x18\t \x03(\t\x12\x12\n\npoint_note\x18\n \x01(\t\x12\x12\n\naudit_mode\x18\x0b \x01(\x08\x12 \n\x08outliers\x18\x0c \x03(\x0b\x32\x0e.SignalOutlier\x12\x15\n\routlier_count\x18\r \x01(\x05\x12\x14\n\x0csample_count\x18\x0e \x01(\x05\x12\x13\n\x0btrend_value\x18\x0f \x01(\x02\x12\x14\n\x0ctrend_margin\x18\x10 \x01(\x02\x12\x16\n\x0ehas_trend_band\x18\x11 \x01(\x08\x12\x11\n\tvalue_min\x18\x12 \x01(\x02\x12\x11\n\tvalue_max\x18\x13 \x01(\x02\x12\x17\n\x0fhas_value_range\x18\x14 \x01(\x08\"\x9a\x01\n\x1bGetLatestLoggerDataResponse\x12 \n\x06points\x18\x01 \x03(\x0b\x32\x10.LoggerDataPoint\x12\x1a\n\x12weightslab_version\x18\x02 \x01(\t\x12!\n\x06\x63urves\x18\x03 \x03(\x0b\x32\x11.SignalCurveIndex\x12\x1a\n\x12\x61pplied_max_points\x18\x04 \x01(\x05\"\x07\n\x05\x45mpty\"/\n\x08NeuronId\x12\x10\n\x08layer_id\x18\x01 \x01(\x05\x12\x11\n\tneuron_id\x18\x02 \x01(\x05\"\x91\x02\n\x0fWeightOperation\x12*\n\x07op_type\x18\x01 \x01(\x0e\x32\x14.WeightOperationTypeH\x00\x88\x01\x01\x12\x15\n\x08layer_id\x18\x02 \x01(\x05H\x01\x88\x01\x01\x12\x1d\n\nneuron_ids\x18\x03 \x03(\x0b\x32\t.NeuronId\x12\x16\n\x0eneurons_to_add\x18\t \x01(\x05\x12 \n\x18zerofy_from_incoming_ids\x18\x0b \x03(\x05\x12\x1c\n\x14zerofy_to_neuron_ids\x18\x0c \x03(\x05\x12+\n\x11zerofy_predicates\x18\r \x03(\x0e\x32\x10.ZerofyPredicateB\n\n\x08_op_typeB\x0b\n\t_layer_id\"_\n\x17WeightsOperationRequest\x12/\n\x10weight_operation\x18\x01 \x01(\x0b\x32\x10.WeightOperationH\x00\x88\x01\x01\x42\x13\n\x11_weight_operation\"<\n\x18WeightsOperationResponse\x12\x0f\n\x07success\x18\x01 \x01(\x08\x12\x0f\n\x07message\x18\x02 \x01(\t\"\xc1\x05\n\x0fHyperParameters\x12\x1c\n\x0f\x65xperiment_name\x18\x01 \x01(\tH\x00\x88\x01\x01\x12!\n\x14training_steps_to_do\x18\x02 \x01(\x05H\x01\x88\x01\x01\x12\x1a\n\rlearning_rate\x18\x03 \x01(\x02H\x02\x88\x01\x01\x12\x17\n\nbatch_size\x18\x04 \x01(\x05H\x03\x88\x01\x01\x12 \n\x13\x66ull_eval_frequency\x18\x05 \x01(\x05H\x04\x88\x01\x01\x12 \n\x13\x63heckpont_frequency\x18\x06 \x01(\x05H\x05\x88\x01\x01\x12\x18\n\x0bis_training\x18\x07 \x01(\x08H\x06\x88\x01\x01\x12\x15\n\x08nb_steps\x18\x08 \x01(\x05H\x07\x88\x01\x01\x12\x19\n\x0c\x61uditor_mode\x18\t \x01(\x08H\x08\x88\x01\x01\x12\x1d\n\x10train_batch_size\x18\n \x01(\x05H\t\x88\x01\x01\x12\x1b\n\x0eval_batch_size\x18\x0b \x01(\x05H\n\x88\x01\x01\x12\x1c\n\x0ftest_batch_size\x18\x0c \x01(\x05H\x0b\x88\x01\x01\x12\x1c\n\x0f\x65valuation_mode\x18\r \x01(\x08H\x0c\x88\x01\x01\x12\x1e\n\x11\x65valuation_config\x18\x0e \x01(\tH\r\x88\x01\x01\x42\x12\n\x10_experiment_nameB\x17\n\x15_training_steps_to_doB\x10\n\x0e_learning_rateB\r\n\x0b_batch_sizeB\x16\n\x14_full_eval_frequencyB\x16\n\x14_checkpont_frequencyB\x0e\n\x0c_is_trainingB\x0b\n\t_nb_stepsB\x0f\n\r_auditor_modeB\x13\n\x11_train_batch_sizeB\x11\n\x0f_val_batch_sizeB\x12\n\x10_test_batch_sizeB\x12\n\x10_evaluation_modeB\x14\n\x12_evaluation_config\",\n\rMetricsStatus\x12\x0c\n\x04name\x18\x01 \x01(\t\x12\r\n\x05value\x18\x02 \x01(\x02\"~\n\rAnnotatStatus\x12\x0c\n\x04name\x18\x01 \x01(\t\x12.\n\x08metadata\x18\x02 \x03(\x0b\x32\x1c.AnnotatStatus.MetadataEntry\x1a/\n\rMetadataEntry\x12\x0b\n\x03key\x18\x01 \x01(\t\x12\r\n\x05value\x18\x02 \x01(\x02:\x02\x38\x01\"\x90\x02\n\x10TrainingStatusEx\x12\x16\n\ttimestamp\x18\x01 \x01(\tH\x00\x88\x01\x01\x12\x1c\n\x0f\x65xperiment_name\x18\x02 \x01(\tH\x01\x88\x01\x01\x12\x16\n\tmodel_age\x18\x03 \x01(\x05H\x02\x88\x01\x01\x12+\n\x0emetrics_status\x18\x04 \x01(\x0b\x32\x0e.MetricsStatusH\x03\x88\x01\x01\x12+\n\x0e\x61nnotat_status\x18\x05 \x01(\x0b\x32\x0e.AnnotatStatusH\x04\x88\x01\x01\x42\x0c\n\n_timestampB\x12\n\x10_experiment_nameB\x0c\n\n_model_ageB\x11\n\x0f_metrics_statusB\x11\n\x0f_annotat_status\"]\n\x15HyperParameterCommand\x12/\n\x10hyper_parameters\x18\x01 \x01(\x0b\x32\x10.HyperParametersH\x00\x88\x01\x01\x42\x13\n\x11_hyper_parameters\">\n\x14\x44\x65nySamplesOperation\x12\x12\n\nsample_ids\x18\x01 \x03(\t\x12\x12\n\naccumulate\x18\x02 \x01(\x08\"0\n\x17LoadCheckpointOperation\x12\x15\n\rcheckpoint_id\x18\x01 \x01(\x05\"b\n\x11PlotNoteOperation\x12\x13\n\x0bmetric_name\x18\x01 \x01(\t\x12\x17\n\x0f\x65xperiment_hash\x18\x02 \x01(\t\x12\x11\n\tmodel_age\x18\x03 \x01(\x05\x12\x0c\n\x04note\x18\x04 \x01(\t\"L\n\x17SaveCheckpointOperation\x12\x19\n\x11save_architecture\x18\x01 \x01(\x08\x12\x16\n\x0esave_optimizer\x18\x02 \x01(\x08\"\x1a\n\x18RestartInstanceOperation\"\x8d\x08\n\x0eTrainerCommand\x12\x1c\n\x14get_hyper_parameters\x18\x04 \x01(\x08\x12\x1e\n\x16get_interactive_layers\x18\x05 \x01(\x08\x12\x1d\n\x10get_data_records\x18\x06 \x01(\tH\x00\x88\x01\x01\x12%\n\x18get_single_layer_info_id\x18\x08 \x01(\x05H\x01\x88\x01\x01\x12;\n\x16hyper_parameter_change\x18\x01 \x01(\x0b\x32\x16.HyperParameterCommandH\x02\x88\x01\x01\x12:\n\x16\x64\x65ny_samples_operation\x18\x07 \x01(\x0b\x32\x15.DenySamplesOperationH\x03\x88\x01\x01\x12?\n\x1b\x64\x65ny_eval_samples_operation\x18\n \x01(\x0b\x32\x15.DenySamplesOperationH\x04\x88\x01\x01\x12@\n\x19load_checkpoint_operation\x18\t \x01(\x0b\x32\x18.LoadCheckpointOperationH\x05\x88\x01\x01\x12\x42\n\x1eremove_from_denylist_operation\x18\x0b \x01(\x0b\x32\x15.DenySamplesOperationH\x06\x88\x01\x01\x12G\n#remove_eval_from_denylist_operation\x18\x0c \x01(\x0b\x32\x15.DenySamplesOperationH\x07\x88\x01\x01\x12\x34\n\x13plot_note_operation\x18\r \x01(\x0b\x32\x12.PlotNoteOperationH\x08\x88\x01\x01\x12@\n\x19save_checkpoint_operation\x18\x0e \x01(\x0b\x32\x18.SaveCheckpointOperationH\t\x88\x01\x01\x12\x39\n\x11restart_operation\x18\x0f \x01(\x0b\x32\x19.RestartInstanceOperationH\n\x88\x01\x01\x42\x13\n\x11_get_data_recordsB\x1b\n\x19_get_single_layer_info_idB\x19\n\x17_hyper_parameter_changeB\x19\n\x17_deny_samples_operationB\x1e\n\x1c_deny_eval_samples_operationB\x1c\n\x1a_load_checkpoint_operationB!\n\x1f_remove_from_denylist_operationB&\n$_remove_eval_from_denylist_operationB\x16\n\x14_plot_note_operationB\x1c\n\x1a_save_checkpoint_operationB\x14\n\x12_restart_operation\"\x9d\x01\n\x12HyperParameterDesc\x12\r\n\x05label\x18\x01 \x01(\t\x12\x0c\n\x04name\x18\x02 \x01(\t\x12\x0c\n\x04type\x18\x03 \x01(\t\x12\x1c\n\x0fnumerical_value\x18\x04 \x01(\x02H\x00\x88\x01\x01\x12\x19\n\x0cstring_value\x18\x05 \x01(\tH\x01\x88\x01\x01\x42\x12\n\x10_numerical_valueB\x0f\n\r_string_value\"\xf2\x02\n\x10NeuronStatistics\x12!\n\tneuron_id\x18\x01 \x01(\x0b\x32\t.NeuronIdH\x00\x88\x01\x01\x12\x17\n\nneuron_age\x18\x02 \x01(\x05H\x01\x88\x01\x01\x12\x1f\n\x12train_trigger_rate\x18\x03 \x01(\x02H\x02\x88\x01\x01\x12\x1e\n\x11\x65val_trigger_rate\x18\x04 \x01(\x02H\x03\x88\x01\x01\x12\x1a\n\rlearning_rate\x18\x07 \x01(\x02H\x04\x88\x01\x01\x12\x36\n\x0bincoming_lr\x18\x08 \x03(\x0b\x32!.NeuronStatistics.IncomingLrEntry\x1a\x31\n\x0fIncomingLrEntry\x12\x0b\n\x03key\x18\x01 \x01(\x05\x12\r\n\x05value\x18\x02 \x01(\x02:\x02\x38\x01\x42\x0c\n\n_neuron_idB\r\n\x0b_neuron_ageB\x15\n\x13_train_trigger_rateB\x14\n\x12_eval_trigger_rateB\x10\n\x0e_learning_rate\"\xf0\x02\n\x13LayerRepresentation\x12\x15\n\x08layer_id\x18\x01 \x01(\x05H\x00\x88\x01\x01\x12\x17\n\nlayer_name\x18\x02 \x01(\tH\x01\x88\x01\x01\x12\x17\n\nlayer_type\x18\x03 \x01(\tH\x02\x88\x01\x01\x12\x1a\n\rneurons_count\x18\x04 \x01(\x05H\x03\x88\x01\x01\x12#\n\x16incoming_neurons_count\x18\x05 \x01(\x05H\x04\x88\x01\x01\x12\x18\n\x0bkernel_size\x18\x06 \x01(\x05H\x05\x88\x01\x01\x12\x13\n\x06stride\x18\x07 \x01(\x05H\x06\x88\x01\x01\x12-\n\x12neurons_statistics\x18\n \x03(\x0b\x32\x11.NeuronStatisticsB\x0b\n\t_layer_idB\r\n\x0b_layer_nameB\r\n\x0b_layer_typeB\x10\n\x0e_neurons_countB\x19\n\x17_incoming_neurons_countB\x0e\n\x0c_kernel_sizeB\t\n\x07_stride\"H\n\x11\x41\x63tivationRequest\x12\x10\n\x08layer_id\x18\x01 \x01(\x05\x12\x11\n\tsample_id\x18\x02 \x01(\t\x12\x0e\n\x06origin\x18\x03 \x01(\t\"H\n\rActivationMap\x12\x11\n\tneuron_id\x18\x01 \x01(\x05\x12\x0e\n\x06values\x18\x02 \x03(\x02\x12\t\n\x01H\x18\x03 \x01(\x05\x12\t\n\x01W\x18\x04 \x01(\x05\"d\n\x12\x41\x63tivationResponse\x12\x12\n\nlayer_type\x18\x01 \x01(\t\x12\x15\n\rneurons_count\x18\x02 \x01(\x05\x12#\n\x0b\x61\x63tivations\x18\x03 \x03(\x0b\x32\x0e.ActivationMap\"\x93\x01\n\tTaskField\x12\x0c\n\x04name\x18\x01 \x01(\t\x12\x15\n\x0b\x66loat_value\x18\x02 \x01(\x02H\x00\x12\x13\n\tint_value\x18\x03 \x01(\x05H\x00\x12\x16\n\x0cstring_value\x18\x04 \x01(\tH\x00\x12\x15\n\x0b\x62ytes_value\x18\x05 \x01(\x0cH\x00\x12\x14\n\nbool_value\x18\x06 \x01(\x08H\x00\x42\x07\n\x05value\"\x87\x03\n\x0eRecordMetadata\x12\x11\n\tsample_id\x18\x01 \x01(\t\x12\x14\n\x0csample_label\x18\x02 \x03(\x05\x12\x19\n\x11sample_prediction\x18\x03 \x03(\x05\x12=\n\x10sample_last_loss\x18\x04 \x03(\x0b\x32#.RecordMetadata.SampleLastLossEntry\x12\x19\n\x11sample_encounters\x18\x05 \x01(\x05\x12\x18\n\x10sample_discarded\x18\x06 \x01(\x08\x12 \n\x0c\x65xtra_fields\x18\x07 \x03(\x0b\x32\n.TaskField\x12\x16\n\x0eprediction_raw\x18\t \x01(\x0c\x12\x11\n\ttask_type\x18\n \x01(\t\x12\x19\n\x11sample_label_text\x18\x0b \x03(\t\x12\x1e\n\x16sample_prediction_text\x18\x0c \x03(\t\x1a\x35\n\x13SampleLastLossEntry\x12\x0b\n\x03key\x18\x01 \x01(\t\x12\r\n\x05value\x18\x02 \x01(\x02:\x02\x38\x01\"\x93\x01\n\x10SampleStatistics\x12\x13\n\x06origin\x18\x06 \x01(\tH\x00\x88\x01\x01\x12\x19\n\x0csample_count\x18\x07 \x01(\x05H\x01\x88\x01\x01\x12\x11\n\ttask_type\x18\t \x01(\t\x12 \n\x07records\x18\x08 \x03(\x0b\x32\x0f.RecordMetadataB\t\n\x07_originB\x0f\n\r_sample_count\"\xe6\x01\n\x0f\x43ommandResponse\x12\x0f\n\x07success\x18\x01 \x01(\x08\x12\x0f\n\x07message\x18\x02 \x01(\t\x12\x33\n\x16hyper_parameters_descs\x18\x03 \x03(\x0b\x32\x13.HyperParameterDesc\x12\x33\n\x15layer_representations\x18\x04 \x03(\x0b\x32\x14.LayerRepresentation\x12\x31\n\x11sample_statistics\x18\x05 \x01(\x0b\x32\x11.SampleStatisticsH\x00\x88\x01\x01\x42\x14\n\x12_sample_statistics\"U\n\rSampleRequest\x12\x16\n\tsample_id\x18\x01 \x01(\tH\x00\x88\x01\x01\x12\x13\n\x06origin\x18\x02 \x01(\tH\x01\x88\x01\x01\x42\x0c\n\n_sample_idB\t\n\x07_origin\"\xad\x02\n\x15SampleRequestResponse\x12\x16\n\tsample_id\x18\x01 \x01(\tH\x00\x88\x01\x01\x12\x13\n\x06origin\x18\x02 \x01(\tH\x01\x88\x01\x01\x12\x12\n\x05label\x18\x03 \x01(\x05H\x02\x88\x01\x01\x12\x11\n\x04\x64\x61ta\x18\x04 \x01(\x0cH\x03\x88\x01\x01\x12\x1a\n\rerror_message\x18\x05 \x01(\tH\x04\x88\x01\x01\x12\x15\n\x08raw_data\x18\x06 \x01(\x0cH\x05\x88\x01\x01\x12\x11\n\x04mask\x18\x07 \x01(\x0cH\x06\x88\x01\x01\x12\x17\n\nprediction\x18\x08 \x01(\x0cH\x07\x88\x01\x01\x42\x0c\n\n_sample_idB\t\n\x07_originB\x08\n\x06_labelB\x07\n\x05_dataB\x10\n\x0e_error_messageB\x0b\n\t_raw_dataB\x07\n\x05_maskB\r\n\x0b_prediction\"\x92\x01\n\x12\x42\x61tchSampleRequest\x12\x12\n\nsample_ids\x18\x01 \x03(\t\x12\x0e\n\x06origin\x18\x02 \x01(\t\x12\x19\n\x0cresize_width\x18\x03 \x01(\x05H\x00\x88\x01\x01\x12\x1a\n\rresize_height\x18\x04 \x01(\x05H\x01\x88\x01\x01\x42\x0f\n\r_resize_widthB\x10\n\x0e_resize_height\">\n\x13\x42\x61tchSampleResponse\x12\'\n\x07samples\x18\x01 \x03(\x0b\x32\x16.SampleRequestResponse\".\n\x0eWeightsRequest\x12\x1c\n\tneuron_id\x18\x01 \x01(\x0b\x32\t.NeuronId\"\x9d\x02\n\x0fWeightsResponse\x12\x1c\n\tneuron_id\x18\x01 \x01(\x0b\x32\t.NeuronId\x12\x17\n\nlayer_name\x18\x02 \x01(\tH\x00\x88\x01\x01\x12\x17\n\nlayer_type\x18\x03 \x01(\tH\x01\x88\x01\x01\x12\x10\n\x08incoming\x18\x04 \x01(\x05\x12\x10\n\x08outgoing\x18\x05 \x01(\x05\x12\x18\n\x0bkernel_size\x18\x06 \x01(\x05H\x02\x88\x01\x01\x12\x0f\n\x07weights\x18\x07 \x03(\x02\x12\x0f\n\x07success\x18\x0b \x01(\x08\x12\x1a\n\rerror_message\x18\x0c \x01(\tH\x03\x88\x01\x01\x42\r\n\x0b_layer_nameB\r\n\x0b_layer_typeB\x0e\n\x0c_kernel_sizeB\x10\n\x0e_error_message\"R\n\x10\x44\x61taQueryRequest\x12\r\n\x05query\x18\x01 \x01(\t\x12\x12\n\naccumulate\x18\x02 \x01(\x08\x12\x1b\n\x13is_natural_language\x18\x03 \x01(\x08\"5\n\x11\x43\x61tegoricalTagDef\x12\x0c\n\x04name\x18\x01 \x01(\t\x12\x12\n\ncategories\x18\x02 \x03(\t\"\xa9\x02\n\x11\x44\x61taQueryResponse\x12\x0f\n\x07success\x18\x01 \x01(\x08\x12\x0f\n\x07message\x18\x02 \x01(\t\x12\x1d\n\x15number_of_all_samples\x18\x03 \x01(\x05\x12%\n\x1dnumber_of_samples_in_the_loop\x18\x04 \x01(\x05\x12#\n\x1bnumber_of_discarded_samples\x18\x05 \x01(\x05\x12\x13\n\x0bunique_tags\x18\x06 \x03(\t\x12+\n\x11\x61gent_intent_type\x18\x07 \x01(\x0e\x32\x10.AgentIntentType\x12\x17\n\x0f\x61nalysis_result\x18\x08 \x01(\t\x12,\n\x10\x63\x61tegorical_tags\x18\t \x03(\x0b\x32\x12.CategoricalTagDef\"\xc2\x01\n\x12\x44\x61taSamplesRequest\x12\x13\n\x0bstart_index\x18\x01 \x01(\x05\x12\x13\n\x0brecords_cnt\x18\x02 \x01(\x05\x12 \n\x18include_transformed_data\x18\x03 \x01(\x08\x12\x18\n\x10include_raw_data\x18\x04 \x01(\x08\x12\x19\n\x11stats_to_retrieve\x18\x05 \x03(\t\x12\x14\n\x0cresize_width\x18\x06 \x01(\x05\x12\x15\n\rresize_height\x18\x07 \x01(\x05\"m\n\x08\x44\x61taStat\x12\x0c\n\x04name\x18\x01 \x01(\t\x12\x0c\n\x04type\x18\x02 \x01(\t\x12\r\n\x05shape\x18\x03 \x03(\x05\x12\r\n\x05value\x18\x04 \x03(\x02\x12\x14\n\x0cvalue_string\x18\x05 \x01(\t\x12\x11\n\tthumbnail\x18\x06 \x01(\x0c\">\n\nDataRecord\x12\x11\n\tsample_id\x18\x01 \x01(\t\x12\x1d\n\ndata_stats\x18\x02 \x03(\x0b\x32\t.DataStat\"\x97\x01\n\x13\x44\x61taSamplesResponse\x12\x0f\n\x07success\x18\x01 \x01(\x08\x12\x0f\n\x07message\x18\x02 \x01(\t\x12!\n\x0c\x64\x61ta_records\x18\x03 \x03(\x0b\x32\x0b.DataRecord\x12\x12\n\nis_subview\x18\x04 \x01(\x08\x12\x12\n\nview_count\x18\x05 \x01(\x03\x12\x13\n\x0btotal_count\x18\x06 \x01(\x03\"C\n\x0fHistogramSubBar\x12\x0e\n\x06origin\x18\x01 \x01(\t\x12\x11\n\tdiscarded\x18\x02 \x01(\x08\x12\r\n\x05\x63ount\x18\x03 \x01(\x03\"h\n\x0cHistogramBin\x12\x0b\n\x03min\x18\x01 \x01(\x01\x12\x0b\n\x03max\x18\x02 \x01(\x01\x12\x0b\n\x03\x61vg\x18\x03 \x01(\x01\x12\r\n\x05\x63ount\x18\x04 \x01(\x03\x12\"\n\x08sub_bars\x18\x05 \x03(\x0b\x32\x10.HistogramSubBar\"[\n\x17\x43\x61tegoricalHistogramBar\x12\r\n\x05label\x18\x01 \x01(\t\x12\r\n\x05\x63ount\x18\x02 \x01(\x03\x12\"\n\x08sub_bars\x18\x03 \x03(\x0b\x32\x10.HistogramSubBar\"4\n\x10HistogramRequest\x12\x0e\n\x06\x63olumn\x18\x01 \x01(\t\x12\x10\n\x08max_bins\x18\x02 \x01(\x05\"\xb2\x01\n\x11HistogramResponse\x12\x0f\n\x07success\x18\x01 \x01(\x08\x12\x0f\n\x07message\x18\x02 \x01(\t\x12\x12\n\ntotal_rows\x18\x03 \x01(\x03\x12\x1b\n\x04\x62ins\x18\x04 \x03(\x0b\x32\r.HistogramBin\x12\x16\n\x0eis_categorical\x18\x05 \x01(\x08\x12\x32\n\x10\x63\x61tegorical_bars\x18\x06 \x03(\x0b\x32\x18.CategoricalHistogramBar\"W\n\x12GetMetaDataRequest\x12\x13\n\x0bstart_index\x18\x01 \x01(\x05\x12\x13\n\x0brecords_cnt\x18\x02 \x01(\x05\x12\x17\n\x0fmodal_sample_id\x18\x03 \x01(\t\"\x99\x01\n\x13GetMetaDataResponse\x12\x0f\n\x07success\x18\x01 \x01(\x08\x12\x0f\n\x07message\x18\x02 \x01(\t\x12\x1a\n\x12\x61ll_metadata_names\x18\x03 \x03(\t\x12!\n\x0cgrid_records\x18\x04 \x03(\x0b\x32\x0b.DataRecord\x12!\n\x0cmodal_record\x18\x05 \x01(\x0b\x32\x0b.DataRecord\"j\n\x12StepSamplesRequest\x12\x13\n\x0bmetric_name\x18\x01 \x01(\t\x12\x17\n\x0f\x65xperiment_hash\x18\x02 \x01(\t\x12\x11\n\tmodel_age\x18\x03 \x01(\x05\x12\x13\n\x0bmax_samples\x18\x04 \x01(\x05\"{\n\x13StepSamplesResponse\x12\x0f\n\x07success\x18\x01 \x01(\x08\x12\x0f\n\x07message\x18\x02 \x01(\t\x12\x12\n\nsample_ids\x18\x03 \x03(\t\x12\x17\n\x0ftotal_available\x18\x04 \x01(\x05\x12\x15\n\rsample_values\x18\x05 \x03(\x02\"Y\n\x1aGetSignalTrajectoryRequest\x12\x13\n\x0bsignal_name\x18\x01 \x01(\t\x12\x12\n\nsample_ids\x18\x02 \x03(\t\x12\x12\n\nmax_points\x18\x03 \x01(\x05\"4\n\x10SignalTrajectory\x12\x11\n\tsample_id\x18\x01 \x01(\t\x12\r\n\x05value\x18\x02 \x03(\x02\"}\n\x1bGetSignalTrajectoryResponse\x12\x0f\n\x07success\x18\x01 \x01(\x08\x12\x0f\n\x07message\x18\x02 \x01(\t\x12\x13\n\x0bsignal_name\x18\x03 \x01(\t\x12\'\n\x0ctrajectories\x18\x04 \x03(\x0b\x32\x11.SignalTrajectory\"Y\n\x11PointCloudRequest\x12\x11\n\tsample_id\x18\x01 \x01(\t\x12\x0e\n\x06origin\x18\x02 \x01(\t\x12\x12\n\nmax_points\x18\x03 \x01(\x05\x12\r\n\x05\x66ield\x18\x04 \x01(\t\"\xbf\x01\n\x0fPointCloudChunk\x12\x0f\n\x07success\x18\x01 \x01(\x08\x12\x0f\n\x07message\x18\x02 \x01(\t\x12\x12\n\nnum_points\x18\x03 \x01(\x05\x12\x14\n\x0cnum_features\x18\x04 \x01(\x05\x12\x10\n\x08pc_range\x18\x05 \x03(\x02\x12\x0c\n\x04\x64\x61ta\x18\x06 \x01(\x0c\x12\x13\n\x0b\x63hunk_index\x18\x07 \x01(\x05\x12\x14\n\x0ctotal_chunks\x18\x08 \x01(\x05\x12\x15\n\rfeature_names\x18\t \x03(\t\"b\n\x0cMediaRequest\x12\x11\n\tsample_id\x18\x01 \x01(\t\x12\x0e\n\x06origin\x18\x02 \x01(\t\x12\x0c\n\x04kind\x18\x03 \x01(\t\x12\x12\n\nmax_frames\x18\x04 \x01(\x05\x12\r\n\x05\x66ield\x18\x05 \x01(\t\"\x92\x02\n\nMediaChunk\x12\x0f\n\x07success\x18\x01 \x01(\x08\x12\x0f\n\x07message\x18\x02 \x01(\t\x12\x11\n\tmime_type\x18\x03 \x01(\t\x12\x13\n\x0b\x66rame_count\x18\x04 \x01(\x05\x12\x0b\n\x03\x66ps\x18\x05 \x01(\x02\x12\x11\n\thas_audio\x18\x06 \x01(\x08\x12\r\n\x05width\x18\x07 \x01(\x05\x12\x0e\n\x06height\x18\x08 \x01(\x05\x12\x13\n\x0btotal_bytes\x18\t \x01(\x05\x12\x18\n\x10\x64uration_seconds\x18\n \x01(\x02\x12\x13\n\x0bsample_rate\x18\x0b \x01(\x05\x12\x0c\n\x04\x64\x61ta\x18\x0c \x01(\x0c\x12\x13\n\x0b\x63hunk_index\x18\r \x01(\x05\x12\x14\n\x0ctotal_chunks\x18\x0e \x01(\x05\"\x8c\x02\n\x10\x44\x61taEditsRequest\x12\x11\n\tstat_name\x18\x01 \x01(\t\x12\x13\n\x0b\x66loat_value\x18\x02 \x01(\x02\x12\x14\n\x0cstring_value\x18\x03 \x01(\t\x12\x12\n\nbool_value\x18\x04 \x01(\x08\x12\x1d\n\x04type\x18\x05 \x01(\x0e\x32\x0f.SampleEditType\x12\x13\n\x0bsamples_ids\x18\x06 \x03(\t\x12\x16\n\x0esample_origins\x18\x07 \x03(\t\x12\x16\n\x0eis_categorical\x18\x08 \x01(\x08\x12\x12\n\ncategories\x18\t \x03(\t\x12\x15\n\rsample_values\x18\n \x03(\x02\x12\x17\n\x0f\x65xperiment_hash\x18\x0b \x01(\t\"5\n\x11\x44\x61taEditsResponse\x12\x0f\n\x07success\x18\x01 \x01(\x08\x12\x0f\n\x07message\x18\x02 \x01(\t\":\n\x12\x44\x61taSplitsResponse\x12\x0f\n\x07success\x18\x01 \x01(\x08\x12\x13\n\x0bsplit_names\x18\x02 \x03(\t\"P\n\x13\x41gentHealthResponse\x12\x11\n\tavailable\x18\x01 \x01(\x08\x12\x0f\n\x07message\x18\x02 \x01(\t\x12\x15\n\rcurrent_model\x18\x03 \x01(\t\"^\n\x16InitializeAgentRequest\x12\x0f\n\x07\x61pi_key\x18\x01 \x01(\t\x12$\n\x08provider\x18\x02 \x01(\x0e\x32\x12.AgentProviderType\x12\r\n\x05model\x18\x03 \x01(\t\";\n\x17InitializeAgentResponse\x12\x0f\n\x07success\x18\x01 \x01(\x08\x12\x0f\n\x07message\x18\x02 \x01(\t\"(\n\x17\x43hangeAgentModelRequest\x12\r\n\x05model\x18\x01 \x01(\t\"<\n\x18\x43hangeAgentModelResponse\x12\x0f\n\x07success\x18\x01 \x01(\x08\x12\x0f\n\x07message\x18\x02 \x01(\t\"\x17\n\x15GetAgentModelsRequest\"a\n\x16GetAgentModelsResponse\x12\x0f\n\x07success\x18\x01 \x01(\x08\x12\x0e\n\x06models\x18\x02 \x03(\t\x12\x0f\n\x07message\x18\x03 \x01(\t\x12\x15\n\rcurrent_model\x18\x04 \x01(\t\"6\n\x12ResetAgentResponse\x12\x0f\n\x07success\x18\x01 \x01(\x08\x12\x0f\n\x07message\x18\x02 \x01(\t\"=\n\x19\x43learAgentHistoryResponse\x12\x0f\n\x07success\x18\x01 \x01(\x08\x12\x0f\n\x07message\x18\x02 \x01(\t\"?\n\x1b\x43ompactAgentHistoryResponse\x12\x0f\n\x07success\x18\x01 \x01(\x08\x12\x0f\n\x07message\x18\x02 \x01(\t\"\xe5\x01\n\x1cGetAgentContextUsageResponse\x12\x0f\n\x07success\x18\x01 \x01(\x08\x12\x0f\n\x07message\x18\x02 \x01(\t\x12\r\n\x05model\x18\x03 \x01(\t\x12\x16\n\x0e\x63ontext_window\x18\x04 \x01(\x03\x12\x14\n\x0cinput_tokens\x18\x05 \x01(\x03\x12\x15\n\routput_tokens\x18\x06 \x01(\x03\x12\x18\n\x10reasoning_tokens\x18\x07 \x01(\x03\x12\x19\n\x11\x63\x61\x63he_read_tokens\x18\x08 \x01(\x03\x12\x1a\n\x12\x63\x61\x63he_write_tokens\x18\t \x01(\x03\"3\n\x18RestoreCheckpointRequest\x12\x17\n\x0f\x65xperiment_hash\x18\x01 \x01(\t\"=\n\x19RestoreCheckpointResponse\x12\x0f\n\x07success\x18\x01 \x01(\x08\x12\x0f\n\x07message\x18\x02 \x01(\t\"\xa8\x01\n\x11\x45xperimentRunInfo\x12\x17\n\x0f\x65xperiment_hash\x18\x01 \x01(\t\x12\x17\n\x0f\x65xperiment_name\x18\x02 \x01(\t\x12\r\n\x05notes\x18\x03 \x01(\t\x12\x0f\n\x07\x63reated\x18\x04 \x01(\t\x12\x11\n\tlast_used\x18\x05 \x01(\t\x12\x1a\n\x12latest_weight_step\x18\x06 \x01(\x05\x12\x12\n\nis_current\x18\x07 \x01(\x08\"\x1b\n\x19ListExperimentRunsRequest\">\n\x1aListExperimentRunsResponse\x12 \n\x04runs\x18\x01 \x03(\x0b\x32\x12.ExperimentRunInfo\"G\n\x1aRenameExperimentRunRequest\x12\x17\n\x0f\x65xperiment_hash\x18\x01 \x01(\t\x12\x10\n\x08new_name\x18\x02 \x01(\t\"?\n\x1bRenameExperimentRunResponse\x12\x0f\n\x07success\x18\x01 \x01(\x08\x12\x0f\n\x07message\x18\x02 \x01(\t\"F\n\x1cSetExperimentRunNotesRequest\x12\x17\n\x0f\x65xperiment_hash\x18\x01 \x01(\t\x12\r\n\x05notes\x18\x02 \x01(\t\"A\n\x1dSetExperimentRunNotesResponse\x12\x0f\n\x07success\x18\x01 \x01(\x08\x12\x0f\n\x07message\x18\x02 \x01(\t\"R\n\x18TriggerEvaluationRequest\x12\x12\n\nsplit_name\x18\x01 \x01(\t\x12\x0c\n\x04tags\x18\x02 \x03(\t\x12\x14\n\x0cuse_full_set\x18\x03 \x01(\x08\"=\n\x19TriggerEvaluationResponse\x12\x0f\n\x07success\x18\x01 \x01(\x08\x12\x0f\n\x07message\x18\x02 \x01(\t\"\x1c\n\x1aGetEvaluationStatusRequest\"\x81\x01\n\x1bGetEvaluationStatusResponse\x12\x0e\n\x06status\x18\x01 \x01(\t\x12\x0f\n\x07\x63urrent\x18\x02 \x01(\x05\x12\r\n\x05total\x18\x03 \x01(\x05\x12\x0f\n\x07message\x18\x04 \x01(\t\x12\r\n\x05\x65rror\x18\x05 \x01(\t\x12\x12\n\nsplit_name\x18\x06 \x01(\t\")\n\x17\x43\x61ncelEvaluationRequest\x12\x0e\n\x06reason\x18\x01 \x01(\t\"<\n\x18\x43\x61ncelEvaluationResponse\x12\x0f\n\x07success\x18\x01 \x01(\x08\x12\x0f\n\x07message\x18\x02 \x01(\t\"7\n\x16RunNotebookCellRequest\x12\x0c\n\x04\x63ode\x18\x01 \x01(\t\x12\x0f\n\x07\x63\x65ll_id\x18\x02 \x01(\t\"2\n\x10NotebookCellDone\x12\x12\n\nexec_count\x18\x01 \x01(\x05\x12\n\n\x02ok\x18\x02 \x01(\x08\"\x1e\n\x1cInterruptNotebookCellRequest\":\n\x1dInterruptNotebookCellResponse\x12\n\n\x02ok\x18\x01 \x01(\x08\x12\r\n\x05\x65rror\x18\x02 \x01(\t\"\xbd\x01\n\x11NotebookCellChunk\x12\x0f\n\x07\x63\x65ll_id\x18\x01 \x01(\t\x12\x10\n\x06stdout\x18\x02 \x01(\tH\x00\x12\x10\n\x06stderr\x18\x03 \x01(\tH\x00\x12\x15\n\x0bresult_text\x18\x04 \x01(\tH\x00\x12\x13\n\timage_png\x18\x05 \x01(\x0cH\x00\x12\x19\n\x0f\x65rror_traceback\x18\x06 \x01(\tH\x00\x12!\n\x04\x64one\x18\x07 \x01(\x0b\x32\x11.NotebookCellDoneH\x00\x42\t\n\x07payload\"S\n\x10NotebookResponse\x12\x12\n\nipynb_json\x18\x01 \x01(\t\x12\x0f\n\x07\x65xisted\x18\x02 \x01(\x08\x12\x0c\n\x04path\x18\x03 \x01(\t\x12\x0c\n\x04name\x18\x04 \x01(\t\"7\n\x13SaveNotebookRequest\x12\x12\n\nipynb_json\x18\x01 \x01(\t\x12\x0c\n\x04name\x18\x02 \x01(\t\"M\n\x14SaveNotebookResponse\x12\n\n\x02ok\x18\x01 \x01(\x08\x12\x0c\n\x04path\x18\x02 \x01(\t\x12\r\n\x05\x65rror\x18\x03 \x01(\t\x12\x0c\n\x04name\x18\x04 \x01(\t\"C\n\x1bGenerateNotebookCodeRequest\x12\x0e\n\x06prompt\x18\x01 \x01(\t\x12\x14\n\x0c\x63ontext_code\x18\x02 \x01(\t\"\\\n\x1cGenerateNotebookCodeResponse\x12\x0c\n\x04\x63ode\x18\x01 \x01(\t\x12\x13\n\x0b\x65xplanation\x18\x02 \x01(\t\x12\n\n\x02ok\x18\x03 \x01(\x08\x12\r\n\x05\x65rror\x18\x04 \x01(\t\"~\n\x18\x45xportAnnotationsRequest\x12\'\n\x06\x66ormat\x18\x01 \x01(\x0e\x32\x17.AnnotationExportFormat\x12\x0e\n\x06origin\x18\x02 \x01(\t\x12\x1b\n\x13include_predictions\x18\x03 \x01(\x08\x12\x0c\n\x04tags\x18\x04 \x03(\t\"\x88\x01\n\x19\x45xportAnnotationsResponse\x12\x0f\n\x07success\x18\x01 \x01(\x08\x12\x0f\n\x07message\x18\x02 \x01(\t\x12\x0f\n\x07payload\x18\x03 \x01(\x0c\x12\x10\n\x08\x66ilename\x18\x04 \x01(\t\x12\x11\n\tmime_type\x18\x05 \x01(\t\x12\x13\n\x0bimage_count\x18\x06 \x01(\x05\"\xf2\x02\n\x11ProjectionRequest\x12\x12\n\nhas_bounds\x18\x01 \x01(\x08\x12\r\n\x05min_x\x18\x02 \x01(\x02\x12\r\n\x05min_y\x18\x03 \x01(\x02\x12\r\n\x05min_z\x18\x04 \x01(\x02\x12\r\n\x05max_x\x18\x05 \x01(\x02\x12\r\n\x05max_y\x18\x06 \x01(\x02\x12\r\n\x05max_z\x18\x07 \x01(\x02\x12\x12\n\nmax_points\x18\x08 \x01(\x05\x12\x14\n\x0c\x63olor_column\x18\t \x01(\t\x12\x0e\n\x06prefix\x18\n \x01(\t\x12\x13\n\x0b\x66ollow_view\x18\x0b \x01(\x08\x12\x1b\n\x13restrict_sample_ids\x18\x0c \x03(\t\x12\x16\n\x0e\x66rustum_planes\x18\r \x03(\x02\x12(\n\nlasso_mode\x18\x0e \x01(\x0e\x32\x14.ProjectionLassoMode\x12\x10\n\x08lasso_px\x18\x0f \x03(\x02\x12\x16\n\x0elasso_viewport\x18\x10 \x03(\x02\x12\x17\n\x0fview_projection\x18\x11 \x03(\x02\"\xa2\x04\n\x12ProjectionResponse\x12\x0f\n\x07success\x18\x01 \x01(\x08\x12\x0f\n\x07message\x18\x02 \x01(\t\x12\x12\n\nsample_ids\x18\x03 \x03(\t\x12\x0e\n\x06\x63oords\x18\x04 \x03(\x02\x12\x12\n\norigin_ids\x18\x05 \x03(\x05\x12\x0f\n\x07origins\x18\x06 \x03(\t\x12\x14\n\x0c\x63olor_values\x18\x07 \x03(\x02\x12\x14\n\x0c\x63olor_labels\x18\x08 \x03(\t\x12\x10\n\x08returned\x18\t \x01(\x03\x12\x15\n\rtotal_in_view\x18\n \x01(\x03\x12\x17\n\x0ftotal_available\x18\x0b \x01(\x03\x12\x14\n\x0c\x65xtent_min_x\x18\x0c \x01(\x02\x12\x14\n\x0c\x65xtent_min_y\x18\r \x01(\x02\x12\x14\n\x0c\x65xtent_min_z\x18\x0e \x01(\x02\x12\x14\n\x0c\x65xtent_max_x\x18\x0f \x01(\x02\x12\x14\n\x0c\x65xtent_max_y\x18\x10 \x01(\x02\x12\x14\n\x0c\x65xtent_max_z\x18\x11 \x01(\x02\x12\x0c\n\x04\x66its\x18\x12 \x01(\x05\x12\x0c\n\x04\x64ims\x18\x13 \x01(\x05\x12\x1a\n\x12\x61vailable_prefixes\x18\x14 \x03(\t\x12\x11\n\tdiscarded\x18\x15 \x03(\x08\x12\x10\n\x08selected\x18\x16 \x01(\x03\x12\x17\n\x0fselection_token\x18\x17 \x01(\t\x12\x14\n\x0cin_selection\x18\x18 \x03(\x08\x12\r\n\x05layer\x18\x19 \x01(\t\x12\x14\n\x0clayer_detail\x18\x1a \x01(\t*d\n\x13WeightOperationType\x12\n\n\x06ZEROFY\x10\x00\x12\x10\n\x0cREINITIALIZE\x10\x01\x12\n\n\x06\x46REEZE\x10\x02\x12\x12\n\x0eREMOVE_NEURONS\x10\t\x12\x0f\n\x0b\x41\x44\x44_NEURONS\x10\n*o\n\x0fZerofyPredicate\x12\x19\n\x15ZEROFY_PREDICATE_NONE\x10\x00\x12 \n\x1cZEROFY_PREDICATE_WITH_FROZEN\x10\x01\x12\x1f\n\x1bZEROFY_PREDICATE_WITH_OLDER\x10\x02*M\n\x0f\x41gentIntentType\x12\x12\n\x0eINTENT_UNKNOWN\x10\x00\x12\x11\n\rINTENT_FILTER\x10\x01\x12\x13\n\x0fINTENT_ANALYSIS\x10\x02*I\n\x0eSampleEditType\x12\x11\n\rEDIT_OVERRIDE\x10\x00\x12\x13\n\x0f\x45\x44IT_ACCUMULATE\x10\x01\x12\x0f\n\x0b\x45\x44IT_REMOVE\x10\x02*C\n\x11\x41gentProviderType\x12\x17\n\x13PROVIDER_OPENROUTER\x10\x00\x12\x15\n\x11PROVIDER_OPENCODE\x10\x01*m\n\x16\x41nnotationExportFormat\x12\x16\n\x12\x45XPORT_FORMAT_CVAT\x10\x00\x12\x1e\n\x1a\x45XPORT_FORMAT_LABEL_STUDIO\x10\x01\x12\x1b\n\x17\x45XPORT_FORMAT_V7_DARWIN\x10\x02*\x85\x01\n\x13ProjectionLassoMode\x12\x19\n\x15PROJECTION_LASSO_NONE\x10\x00\x12\x1c\n\x18PROJECTION_LASSO_REPLACE\x10\x01\x12\x18\n\x14PROJECTION_LASSO_ADD\x10\x02\x12\x1b\n\x17PROJECTION_LASSO_REFINE\x10\x03\x32\xb8\x13\n\x11\x45xperimentService\x12P\n\x13GetLatestLoggerData\x12\x1b.GetLatestLoggerDataRequest\x1a\x1c.GetLatestLoggerDataResponse\x12\x36\n\x11\x45xperimentCommand\x12\x0f.TrainerCommand\x1a\x10.CommandResponse\x12H\n\x11ManipulateWeights\x12\x18.WeightsOperationRequest\x1a\x19.WeightsOperationResponse\x12/\n\nGetWeights\x12\x0f.WeightsRequest\x1a\x10.WeightsResponse\x12\x39\n\x0eGetActivations\x12\x12.ActivationRequest\x1a\x13.ActivationResponse\x12\x37\n\nGetSamples\x12\x13.BatchSampleRequest\x1a\x14.BatchSampleResponse\x12\x37\n\x0e\x41pplyDataQuery\x12\x11.DataQueryRequest\x1a\x12.DataQueryResponse\x12;\n\x0eGetDataSamples\x12\x13.DataSamplesRequest\x1a\x14.DataSamplesResponse\x12\x35\n\x0cGetHistogram\x12\x11.HistogramRequest\x1a\x12.HistogramResponse\x12\x38\n\x0bGetMetaData\x12\x13.GetMetaDataRequest\x1a\x14.GetMetaDataResponse\x12P\n\x13GetSignalTrajectory\x12\x1b.GetSignalTrajectoryRequest\x1a\x1c.GetSignalTrajectoryResponse\x12;\n\x0eGetStepSamples\x12\x13.StepSamplesRequest\x1a\x14.StepSamplesResponse\x12\x37\n\rGetPointCloud\x12\x12.PointCloudRequest\x1a\x10.PointCloudChunk0\x01\x12(\n\x08GetMedia\x12\r.MediaRequest\x1a\x0b.MediaChunk0\x01\x12\x37\n\x0e\x45\x64itDataSample\x12\x11.DataEditsRequest\x1a\x12.DataEditsResponse\x12,\n\rGetDataSplits\x12\x06.Empty\x1a\x13.DataSplitsResponse\x12\x30\n\x10\x43heckAgentHealth\x12\x06.Empty\x1a\x14.AgentHealthResponse\x12\x44\n\x0fInitializeAgent\x12\x17.InitializeAgentRequest\x1a\x18.InitializeAgentResponse\x12G\n\x10\x43hangeAgentModel\x12\x18.ChangeAgentModelRequest\x1a\x19.ChangeAgentModelResponse\x12\x41\n\x0eGetAgentModels\x12\x16.GetAgentModelsRequest\x1a\x17.GetAgentModelsResponse\x12)\n\nResetAgent\x12\x06.Empty\x1a\x13.ResetAgentResponse\x12\x37\n\x11\x43learAgentHistory\x12\x06.Empty\x1a\x1a.ClearAgentHistoryResponse\x12;\n\x13\x43ompactAgentHistory\x12\x06.Empty\x1a\x1c.CompactAgentHistoryResponse\x12=\n\x14GetAgentContextUsage\x12\x06.Empty\x1a\x1d.GetAgentContextUsageResponse\x12@\n\x0fRunNotebookCell\x12\x17.RunNotebookCellRequest\x1a\x12.NotebookCellChunk0\x01\x12V\n\x15InterruptNotebookCell\x12\x1d.InterruptNotebookCellRequest\x1a\x1e.InterruptNotebookCellResponse\x12(\n\x0bGetNotebook\x12\x06.Empty\x1a\x11.NotebookResponse\x12;\n\x0cSaveNotebook\x12\x14.SaveNotebookRequest\x1a\x15.SaveNotebookResponse\x12S\n\x14GenerateNotebookCode\x12\x1c.GenerateNotebookCodeRequest\x1a\x1d.GenerateNotebookCodeResponse\x12J\n\x11RestoreCheckpoint\x12\x19.RestoreCheckpointRequest\x1a\x1a.RestoreCheckpointResponse\x12M\n\x12ListExperimentRuns\x12\x1a.ListExperimentRunsRequest\x1a\x1b.ListExperimentRunsResponse\x12P\n\x13RenameExperimentRun\x12\x1b.RenameExperimentRunRequest\x1a\x1c.RenameExperimentRunResponse\x12V\n\x15SetExperimentRunNotes\x12\x1d.SetExperimentRunNotesRequest\x1a\x1e.SetExperimentRunNotesResponse\x12J\n\x11TriggerEvaluation\x12\x19.TriggerEvaluationRequest\x1a\x1a.TriggerEvaluationResponse\x12P\n\x13GetEvaluationStatus\x12\x1b.GetEvaluationStatusRequest\x1a\x1c.GetEvaluationStatusResponse\x12G\n\x10\x43\x61ncelEvaluation\x12\x18.CancelEvaluationRequest\x1a\x19.CancelEvaluationResponse\x12J\n\x11\x45xportAnnotations\x12\x19.ExportAnnotationsRequest\x1a\x1a.ExportAnnotationsResponse\x12\x38\n\rGetProjection\x12\x12.ProjectionRequest\x1a\x13.ProjectionResponseb\x06proto3') _globals = globals() _builder.BuildMessageAndEnumDescriptors(DESCRIPTOR, _globals) @@ -37,18 +37,20 @@ _globals['_NEURONSTATISTICS_INCOMINGLRENTRY']._serialized_options = b'8\001' _globals['_RECORDMETADATA_SAMPLELASTLOSSENTRY']._loaded_options = None _globals['_RECORDMETADATA_SAMPLELASTLOSSENTRY']._serialized_options = b'8\001' - _globals['_WEIGHTOPERATIONTYPE']._serialized_start=13491 - _globals['_WEIGHTOPERATIONTYPE']._serialized_end=13591 - _globals['_ZEROFYPREDICATE']._serialized_start=13593 - _globals['_ZEROFYPREDICATE']._serialized_end=13704 - _globals['_AGENTINTENTTYPE']._serialized_start=13706 - _globals['_AGENTINTENTTYPE']._serialized_end=13783 - _globals['_SAMPLEEDITTYPE']._serialized_start=13785 - _globals['_SAMPLEEDITTYPE']._serialized_end=13858 - _globals['_AGENTPROVIDERTYPE']._serialized_start=13860 - _globals['_AGENTPROVIDERTYPE']._serialized_end=13927 - _globals['_ANNOTATIONEXPORTFORMAT']._serialized_start=13929 - _globals['_ANNOTATIONEXPORTFORMAT']._serialized_end=14038 + _globals['_WEIGHTOPERATIONTYPE']._serialized_start=14459 + _globals['_WEIGHTOPERATIONTYPE']._serialized_end=14559 + _globals['_ZEROFYPREDICATE']._serialized_start=14561 + _globals['_ZEROFYPREDICATE']._serialized_end=14672 + _globals['_AGENTINTENTTYPE']._serialized_start=14674 + _globals['_AGENTINTENTTYPE']._serialized_end=14751 + _globals['_SAMPLEEDITTYPE']._serialized_start=14753 + _globals['_SAMPLEEDITTYPE']._serialized_end=14826 + _globals['_AGENTPROVIDERTYPE']._serialized_start=14828 + _globals['_AGENTPROVIDERTYPE']._serialized_end=14895 + _globals['_ANNOTATIONEXPORTFORMAT']._serialized_start=14897 + _globals['_ANNOTATIONEXPORTFORMAT']._serialized_end=15006 + _globals['_PROJECTIONLASSOMODE']._serialized_start=15009 + _globals['_PROJECTIONLASSOMODE']._serialized_end=15142 _globals['_GETLATESTLOGGERDATAREQUEST']._serialized_start=46 _globals['_GETLATESTLOGGERDATAREQUEST']._serialized_end=303 _globals['_SIGNALCURVEINDEX']._serialized_start=306 @@ -182,81 +184,85 @@ _globals['_DATASPLITSRESPONSE']._serialized_start=10479 _globals['_DATASPLITSRESPONSE']._serialized_end=10537 _globals['_AGENTHEALTHRESPONSE']._serialized_start=10539 - _globals['_AGENTHEALTHRESPONSE']._serialized_end=10596 - _globals['_INITIALIZEAGENTREQUEST']._serialized_start=10598 - _globals['_INITIALIZEAGENTREQUEST']._serialized_end=10692 - _globals['_INITIALIZEAGENTRESPONSE']._serialized_start=10694 - _globals['_INITIALIZEAGENTRESPONSE']._serialized_end=10753 - _globals['_CHANGEAGENTMODELREQUEST']._serialized_start=10755 - _globals['_CHANGEAGENTMODELREQUEST']._serialized_end=10795 - _globals['_CHANGEAGENTMODELRESPONSE']._serialized_start=10797 - _globals['_CHANGEAGENTMODELRESPONSE']._serialized_end=10857 - _globals['_GETAGENTMODELSREQUEST']._serialized_start=10859 - _globals['_GETAGENTMODELSREQUEST']._serialized_end=10882 - _globals['_GETAGENTMODELSRESPONSE']._serialized_start=10884 - _globals['_GETAGENTMODELSRESPONSE']._serialized_end=10958 - _globals['_RESETAGENTRESPONSE']._serialized_start=10960 - _globals['_RESETAGENTRESPONSE']._serialized_end=11014 - _globals['_CLEARAGENTHISTORYRESPONSE']._serialized_start=11016 - _globals['_CLEARAGENTHISTORYRESPONSE']._serialized_end=11077 - _globals['_COMPACTAGENTHISTORYRESPONSE']._serialized_start=11079 - _globals['_COMPACTAGENTHISTORYRESPONSE']._serialized_end=11142 - _globals['_GETAGENTCONTEXTUSAGERESPONSE']._serialized_start=11145 - _globals['_GETAGENTCONTEXTUSAGERESPONSE']._serialized_end=11374 - _globals['_RESTORECHECKPOINTREQUEST']._serialized_start=11376 - _globals['_RESTORECHECKPOINTREQUEST']._serialized_end=11427 - _globals['_RESTORECHECKPOINTRESPONSE']._serialized_start=11429 - _globals['_RESTORECHECKPOINTRESPONSE']._serialized_end=11490 - _globals['_EXPERIMENTRUNINFO']._serialized_start=11493 - _globals['_EXPERIMENTRUNINFO']._serialized_end=11661 - _globals['_LISTEXPERIMENTRUNSREQUEST']._serialized_start=11663 - _globals['_LISTEXPERIMENTRUNSREQUEST']._serialized_end=11690 - _globals['_LISTEXPERIMENTRUNSRESPONSE']._serialized_start=11692 - _globals['_LISTEXPERIMENTRUNSRESPONSE']._serialized_end=11754 - _globals['_RENAMEEXPERIMENTRUNREQUEST']._serialized_start=11756 - _globals['_RENAMEEXPERIMENTRUNREQUEST']._serialized_end=11827 - _globals['_RENAMEEXPERIMENTRUNRESPONSE']._serialized_start=11829 - _globals['_RENAMEEXPERIMENTRUNRESPONSE']._serialized_end=11892 - _globals['_SETEXPERIMENTRUNNOTESREQUEST']._serialized_start=11894 - _globals['_SETEXPERIMENTRUNNOTESREQUEST']._serialized_end=11964 - _globals['_SETEXPERIMENTRUNNOTESRESPONSE']._serialized_start=11966 - _globals['_SETEXPERIMENTRUNNOTESRESPONSE']._serialized_end=12031 - _globals['_TRIGGEREVALUATIONREQUEST']._serialized_start=12033 - _globals['_TRIGGEREVALUATIONREQUEST']._serialized_end=12115 - _globals['_TRIGGEREVALUATIONRESPONSE']._serialized_start=12117 - _globals['_TRIGGEREVALUATIONRESPONSE']._serialized_end=12178 - _globals['_GETEVALUATIONSTATUSREQUEST']._serialized_start=12180 - _globals['_GETEVALUATIONSTATUSREQUEST']._serialized_end=12208 - _globals['_GETEVALUATIONSTATUSRESPONSE']._serialized_start=12211 - _globals['_GETEVALUATIONSTATUSRESPONSE']._serialized_end=12340 - _globals['_CANCELEVALUATIONREQUEST']._serialized_start=12342 - _globals['_CANCELEVALUATIONREQUEST']._serialized_end=12383 - _globals['_CANCELEVALUATIONRESPONSE']._serialized_start=12385 - _globals['_CANCELEVALUATIONRESPONSE']._serialized_end=12445 - _globals['_RUNNOTEBOOKCELLREQUEST']._serialized_start=12447 - _globals['_RUNNOTEBOOKCELLREQUEST']._serialized_end=12502 - _globals['_NOTEBOOKCELLDONE']._serialized_start=12504 - _globals['_NOTEBOOKCELLDONE']._serialized_end=12554 - _globals['_INTERRUPTNOTEBOOKCELLREQUEST']._serialized_start=12556 - _globals['_INTERRUPTNOTEBOOKCELLREQUEST']._serialized_end=12586 - _globals['_INTERRUPTNOTEBOOKCELLRESPONSE']._serialized_start=12588 - _globals['_INTERRUPTNOTEBOOKCELLRESPONSE']._serialized_end=12646 - _globals['_NOTEBOOKCELLCHUNK']._serialized_start=12649 - _globals['_NOTEBOOKCELLCHUNK']._serialized_end=12838 - _globals['_NOTEBOOKRESPONSE']._serialized_start=12840 - _globals['_NOTEBOOKRESPONSE']._serialized_end=12923 - _globals['_SAVENOTEBOOKREQUEST']._serialized_start=12925 - _globals['_SAVENOTEBOOKREQUEST']._serialized_end=12980 - _globals['_SAVENOTEBOOKRESPONSE']._serialized_start=12982 - _globals['_SAVENOTEBOOKRESPONSE']._serialized_end=13059 - _globals['_GENERATENOTEBOOKCODEREQUEST']._serialized_start=13061 - _globals['_GENERATENOTEBOOKCODEREQUEST']._serialized_end=13128 - _globals['_GENERATENOTEBOOKCODERESPONSE']._serialized_start=13130 - _globals['_GENERATENOTEBOOKCODERESPONSE']._serialized_end=13222 - _globals['_EXPORTANNOTATIONSREQUEST']._serialized_start=13224 - _globals['_EXPORTANNOTATIONSREQUEST']._serialized_end=13350 - _globals['_EXPORTANNOTATIONSRESPONSE']._serialized_start=13353 - _globals['_EXPORTANNOTATIONSRESPONSE']._serialized_end=13489 - _globals['_EXPERIMENTSERVICE']._serialized_start=14041 - _globals['_EXPERIMENTSERVICE']._serialized_end=16471 + _globals['_AGENTHEALTHRESPONSE']._serialized_end=10619 + _globals['_INITIALIZEAGENTREQUEST']._serialized_start=10621 + _globals['_INITIALIZEAGENTREQUEST']._serialized_end=10715 + _globals['_INITIALIZEAGENTRESPONSE']._serialized_start=10717 + _globals['_INITIALIZEAGENTRESPONSE']._serialized_end=10776 + _globals['_CHANGEAGENTMODELREQUEST']._serialized_start=10778 + _globals['_CHANGEAGENTMODELREQUEST']._serialized_end=10818 + _globals['_CHANGEAGENTMODELRESPONSE']._serialized_start=10820 + _globals['_CHANGEAGENTMODELRESPONSE']._serialized_end=10880 + _globals['_GETAGENTMODELSREQUEST']._serialized_start=10882 + _globals['_GETAGENTMODELSREQUEST']._serialized_end=10905 + _globals['_GETAGENTMODELSRESPONSE']._serialized_start=10907 + _globals['_GETAGENTMODELSRESPONSE']._serialized_end=11004 + _globals['_RESETAGENTRESPONSE']._serialized_start=11006 + _globals['_RESETAGENTRESPONSE']._serialized_end=11060 + _globals['_CLEARAGENTHISTORYRESPONSE']._serialized_start=11062 + _globals['_CLEARAGENTHISTORYRESPONSE']._serialized_end=11123 + _globals['_COMPACTAGENTHISTORYRESPONSE']._serialized_start=11125 + _globals['_COMPACTAGENTHISTORYRESPONSE']._serialized_end=11188 + _globals['_GETAGENTCONTEXTUSAGERESPONSE']._serialized_start=11191 + _globals['_GETAGENTCONTEXTUSAGERESPONSE']._serialized_end=11420 + _globals['_RESTORECHECKPOINTREQUEST']._serialized_start=11422 + _globals['_RESTORECHECKPOINTREQUEST']._serialized_end=11473 + _globals['_RESTORECHECKPOINTRESPONSE']._serialized_start=11475 + _globals['_RESTORECHECKPOINTRESPONSE']._serialized_end=11536 + _globals['_EXPERIMENTRUNINFO']._serialized_start=11539 + _globals['_EXPERIMENTRUNINFO']._serialized_end=11707 + _globals['_LISTEXPERIMENTRUNSREQUEST']._serialized_start=11709 + _globals['_LISTEXPERIMENTRUNSREQUEST']._serialized_end=11736 + _globals['_LISTEXPERIMENTRUNSRESPONSE']._serialized_start=11738 + _globals['_LISTEXPERIMENTRUNSRESPONSE']._serialized_end=11800 + _globals['_RENAMEEXPERIMENTRUNREQUEST']._serialized_start=11802 + _globals['_RENAMEEXPERIMENTRUNREQUEST']._serialized_end=11873 + _globals['_RENAMEEXPERIMENTRUNRESPONSE']._serialized_start=11875 + _globals['_RENAMEEXPERIMENTRUNRESPONSE']._serialized_end=11938 + _globals['_SETEXPERIMENTRUNNOTESREQUEST']._serialized_start=11940 + _globals['_SETEXPERIMENTRUNNOTESREQUEST']._serialized_end=12010 + _globals['_SETEXPERIMENTRUNNOTESRESPONSE']._serialized_start=12012 + _globals['_SETEXPERIMENTRUNNOTESRESPONSE']._serialized_end=12077 + _globals['_TRIGGEREVALUATIONREQUEST']._serialized_start=12079 + _globals['_TRIGGEREVALUATIONREQUEST']._serialized_end=12161 + _globals['_TRIGGEREVALUATIONRESPONSE']._serialized_start=12163 + _globals['_TRIGGEREVALUATIONRESPONSE']._serialized_end=12224 + _globals['_GETEVALUATIONSTATUSREQUEST']._serialized_start=12226 + _globals['_GETEVALUATIONSTATUSREQUEST']._serialized_end=12254 + _globals['_GETEVALUATIONSTATUSRESPONSE']._serialized_start=12257 + _globals['_GETEVALUATIONSTATUSRESPONSE']._serialized_end=12386 + _globals['_CANCELEVALUATIONREQUEST']._serialized_start=12388 + _globals['_CANCELEVALUATIONREQUEST']._serialized_end=12429 + _globals['_CANCELEVALUATIONRESPONSE']._serialized_start=12431 + _globals['_CANCELEVALUATIONRESPONSE']._serialized_end=12491 + _globals['_RUNNOTEBOOKCELLREQUEST']._serialized_start=12493 + _globals['_RUNNOTEBOOKCELLREQUEST']._serialized_end=12548 + _globals['_NOTEBOOKCELLDONE']._serialized_start=12550 + _globals['_NOTEBOOKCELLDONE']._serialized_end=12600 + _globals['_INTERRUPTNOTEBOOKCELLREQUEST']._serialized_start=12602 + _globals['_INTERRUPTNOTEBOOKCELLREQUEST']._serialized_end=12632 + _globals['_INTERRUPTNOTEBOOKCELLRESPONSE']._serialized_start=12634 + _globals['_INTERRUPTNOTEBOOKCELLRESPONSE']._serialized_end=12692 + _globals['_NOTEBOOKCELLCHUNK']._serialized_start=12695 + _globals['_NOTEBOOKCELLCHUNK']._serialized_end=12884 + _globals['_NOTEBOOKRESPONSE']._serialized_start=12886 + _globals['_NOTEBOOKRESPONSE']._serialized_end=12969 + _globals['_SAVENOTEBOOKREQUEST']._serialized_start=12971 + _globals['_SAVENOTEBOOKREQUEST']._serialized_end=13026 + _globals['_SAVENOTEBOOKRESPONSE']._serialized_start=13028 + _globals['_SAVENOTEBOOKRESPONSE']._serialized_end=13105 + _globals['_GENERATENOTEBOOKCODEREQUEST']._serialized_start=13107 + _globals['_GENERATENOTEBOOKCODEREQUEST']._serialized_end=13174 + _globals['_GENERATENOTEBOOKCODERESPONSE']._serialized_start=13176 + _globals['_GENERATENOTEBOOKCODERESPONSE']._serialized_end=13268 + _globals['_EXPORTANNOTATIONSREQUEST']._serialized_start=13270 + _globals['_EXPORTANNOTATIONSREQUEST']._serialized_end=13396 + _globals['_EXPORTANNOTATIONSRESPONSE']._serialized_start=13399 + _globals['_EXPORTANNOTATIONSRESPONSE']._serialized_end=13535 + _globals['_PROJECTIONREQUEST']._serialized_start=13538 + _globals['_PROJECTIONREQUEST']._serialized_end=13908 + _globals['_PROJECTIONRESPONSE']._serialized_start=13911 + _globals['_PROJECTIONRESPONSE']._serialized_end=14457 + _globals['_EXPERIMENTSERVICE']._serialized_start=15145 + _globals['_EXPERIMENTSERVICE']._serialized_end=17633 # @@protoc_insertion_point(module_scope) diff --git a/weightslab/proto/experiment_service_pb2_grpc.py b/weightslab/proto/experiment_service_pb2_grpc.py index dd5d1209..860f113b 100644 --- a/weightslab/proto/experiment_service_pb2_grpc.py +++ b/weightslab/proto/experiment_service_pb2_grpc.py @@ -219,6 +219,11 @@ def __init__(self, channel): request_serializer=weightslab_dot_proto_dot_experiment__service__pb2.ExportAnnotationsRequest.SerializeToString, response_deserializer=weightslab_dot_proto_dot_experiment__service__pb2.ExportAnnotationsResponse.FromString, _registered_method=True) + self.GetProjection = channel.unary_unary( + '/ExperimentService/GetProjection', + request_serializer=weightslab_dot_proto_dot_experiment__service__pb2.ProjectionRequest.SerializeToString, + response_deserializer=weightslab_dot_proto_dot_experiment__service__pb2.ProjectionResponse.FromString, + _registered_method=True) class ExperimentServiceServicer(object): @@ -494,6 +499,17 @@ def ExportAnnotations(self, request, context): context.set_details('Method not implemented!') raise NotImplementedError('Method not implemented!') + def GetProjection(self, request, context): + """The live parametric-UMAP projection of the model's representation, as a + 3-D point per sample. Level-of-detail is server-side and driven by the + client's current view box: narrowing the box as the user zooms in buys + finer detail for the SAME point budget, instead of ever shipping the whole + dataset. See ProjectionRequest. + """ + context.set_code(grpc.StatusCode.UNIMPLEMENTED) + context.set_details('Method not implemented!') + raise NotImplementedError('Method not implemented!') + def add_ExperimentServiceServicer_to_server(servicer, server): rpc_method_handlers = { @@ -682,6 +698,11 @@ def add_ExperimentServiceServicer_to_server(servicer, server): request_deserializer=weightslab_dot_proto_dot_experiment__service__pb2.ExportAnnotationsRequest.FromString, response_serializer=weightslab_dot_proto_dot_experiment__service__pb2.ExportAnnotationsResponse.SerializeToString, ), + 'GetProjection': grpc.unary_unary_rpc_method_handler( + servicer.GetProjection, + request_deserializer=weightslab_dot_proto_dot_experiment__service__pb2.ProjectionRequest.FromString, + response_serializer=weightslab_dot_proto_dot_experiment__service__pb2.ProjectionResponse.SerializeToString, + ), } generic_handler = grpc.method_handlers_generic_handler( 'ExperimentService', rpc_method_handlers) @@ -1691,3 +1712,30 @@ def ExportAnnotations(request, timeout, metadata, _registered_method=True) + + @staticmethod + def GetProjection(request, + target, + options=(), + channel_credentials=None, + call_credentials=None, + insecure=False, + compression=None, + wait_for_ready=None, + timeout=None, + metadata=None): + return grpc.experimental.unary_unary( + request, + target, + '/ExperimentService/GetProjection', + weightslab_dot_proto_dot_experiment__service__pb2.ProjectionRequest.SerializeToString, + weightslab_dot_proto_dot_experiment__service__pb2.ProjectionResponse.FromString, + options, + channel_credentials, + insecure, + call_credentials, + compression, + wait_for_ready, + timeout, + metadata, + _registered_method=True) diff --git a/weightslab/src.py b/weightslab/src.py index b20a21dc..2b76aa93 100644 --- a/weightslab/src.py +++ b/weightslab/src.py @@ -38,6 +38,15 @@ from weightslab.backend import ledgers from weightslab.backend.ledgers import register_signal from weightslab.components.global_monitoring import pause_controller as pause_ctrl, get_active_origin +from weightslab.projection import ( + attach_projection as _attach_projection, + detach_projection as _detach_projection, + observe_batch as _projection_observe_batch, + # Re-exported: `wl.project_dataset` / `wl.save_projection_coords` resolve + # through this module (see _LAZY_EXPORTS in weightslab/__init__.py). + project_dataset as project_dataset, + save_projection_coords as save_projection_coords, +) def _rebind_caller_local(original_obj: Any, new_obj: Any) -> None: @@ -883,6 +892,12 @@ def wrappered_fwd(original_forward, kwargs, reg_name, *a, **kw): # User parameters batch_ids = wl_kw.get('batch_ids') + # Snapshot the SAMPLE-level ids now. The per_instance branch below rebinds + # batch_ids to one entry per annotation (detection/segmentation), which is + # the right thing for save_instance_signals but would silently break the + # projection's positional feature<->id pairing, since the forward hook saw + # one feature row per SAMPLE. + sample_batch_ids = batch_ids group_ids = wl_kw.get('group_id') batch_scalar = wl_kw.get('signals') preds = wl_kw.get('preds') @@ -973,6 +988,13 @@ def wrappered_fwd(original_forward, kwargs, reg_name, *a, **kw): else: _log_signal(scalar, batch_scalar, reg_name, step=step, **kwargs) + # Live 3-D projection. The forward hook captured this batch's penultimate + # features but cannot see sample ids; we have the ids but not the + # activations. Both describe the same batch in the same order, so this is + # where they get paired -- the same arrangement ctx.logits already uses. + # No-op (one attribute read) unless a projection is attached. + _projection_observe_batch(sample_batch_ids, step) + # CHECK FOR SUBSCRIBERS (Dynamic Signals) # Allows @wl.signal(subscribe_to="metric_name") dynamic_updates = {} @@ -1254,6 +1276,12 @@ def watch_or_edit(obj: Callable, obj_name: str = None, flag: str = None, **kwarg if key in kwargs } + # Live 3-D projection options, popped for the same reason: it is a hook + # ON a wrapped model, not part of what wrapping means. False opts out; + # a dict steers layer/dim/cadence; a string is shorthand for + # {'layer': ...}; absent (or True) means "the env var decides". + projection_opts = kwargs.pop('projection', None) + # Now construct the wrapper and let it register into the ledger. wrapper = ModelInterface(obj, **kwargs) if forced_model_wrapping or _model == None else _model @@ -1281,6 +1309,33 @@ def watch_or_edit(obj: Callable, obj_name: str = None, flag: str = None, **kwarg # without these curves. logger.warning(f"Could not install model signal tracking: {exc}") + # Live parametric-UMAP projection of the penultimate features. On by + # default (WEIGHTSLAB_PROJECTION=0 removes the hook entirely); pass + # projection={'layer': 'backbone.fc', ...} to steer it, or + # projection=False to opt this one model out. Same failure contract as + # the signal tracking above: a projection that cannot install is a lost + # view, never a failed run. + if projection_opts is False: + # One projection per process, and this is THE model now: a hook + # left from a previous wrap would keep fitting on a model nobody + # trains, and a checkpoint restore parked for it would be applied + # to the next one that opts in. + _detach_projection() + from weightslab.projection import clear_pending_restore + clear_pending_restore() + logger.info("Live projection off for this model (projection=False)") + else: + try: + if isinstance(projection_opts, dict): + opts = dict(projection_opts) + elif isinstance(projection_opts, str): + opts = {'layer': projection_opts} + else: + opts = {} + _attach_projection(obj, layer=opts.pop('layer', None), **opts) + except Exception as exc: + logger.warning(f"Could not install live projection: {exc}") + # Prefer returning the proxy (if one exists) so external callers hold # a stable reference that will see updates. If no proxy was # obtainable, return the wrapper itself. @@ -1623,6 +1678,23 @@ def start_training(timeout: int = None) -> None: pause_ctrl.resume(force=not ledgers.list_dataloaders()) +def pause_training() -> None: + """Pause training from your own code — the Studio's pause button, as a call. + + The counterpart of :func:`start_training`. Takes effect at the next + ``with wl.guard_training_context:``, which blocks until training is resumed + (the Play button in the Studio header, or :func:`start_training`). The step + in flight is not interrupted. While paused, the gRPC server, the Studio and + the notebook keep running against the live model and data — which is the + point: pause once the run has converged, then inspect it. + + Example: + >>> if converged: + ... wl.pause_training() # the next training step waits for Play + """ + pause_ctrl.pause() + + def _register_pid_with_ui_server() -> None: """Tell the ``weightslab start`` UI server that this training process exists, so stopping that workspace (Ctrl+C, closing the terminal) stops @@ -2405,6 +2477,7 @@ def save_signals( step: int | None = None, log: bool = False, _react: bool = True, + _seen: bool = True, ): """Save **per-sample** statistics to the tracked dataset. @@ -2421,12 +2494,21 @@ def save_signals( Shapes / formats: - ``signals`` values: shape ``(B,)`` (one scalar per sample) or ``(B, ...)`` — extra dims are mean-reduced to ``(B,)`` before storage. + Accepted types are ``th.Tensor``, ``np.ndarray``, ``list`` and scalar + numbers; anything else raises ``TypeError`` rather than creating a + column that is never written to. - ``batch_ids``: length ``B``; the i-th entry names the sample the i-th row of every other array belongs to. Coerced to ``str`` internally. - ``preds`` / ``preds_raw`` / ``targets``: array of length ``B`` (or a ``dict`` of such arrays, or a ``list`` of length ``B`` for inhomogeneous per-sample shapes). 1-D arrays get a trailing axis. + Raises: + TypeError: a value in *signals* is of a type that cannot be turned into + one number per sample. This is deliberately loud: such a value used + to be silently dropped, leaving a column in the dataframe that stayed + empty forever with nothing anywhere saying why. + Args: signals (dict | th.Tensor): ``{name: values}`` where ``values`` is a length-``B`` tensor/array (or a bare tensor → stored as @@ -2532,39 +2614,83 @@ def normalize(x): return [to_numpy(t) for t in x] if isinstance(x, th.Tensor): return to_numpy(x) + # A bare numpy array is as ordinary an input as a tensor -- half of + # WeightsLab's own signal helpers compute in numpy. Returning None for + # it (as this used to) created the column and then wrote nothing into + # it, which surfaces as a metric that silently stays empty forever. + if isinstance(x, np.ndarray): + return to_numpy(x) + # Scalars, so a caller who computed one value per sample in a Python + # list-comprehension is not silently dropped either. + if isinstance(x, (int, float, np.number)): + return to_numpy(np.asarray(x)) return None preds_np = normalize(preds) preds_raw_np = normalize(preds_raw) target_np = normalize(targets) + def _reduce_signal(arr): + """Collapse a per-sample signal to one value per sample.""" + if isinstance(arr, np.ndarray) and arr.ndim > 1: + return arr.mean(axis=tuple(range(1, arr.ndim))) + return arr + + def _coerce_signal(name, value): + """Normalize one signal value, and REFUSE to write nothing. + + Anything ``normalize`` cannot make sense of used to come back as + ``None``, which ``enqueue_batch`` happily accepted: the column appeared + in the dataframe, held no values, and nothing anywhere said why. That + is the worst possible failure for an observability feature -- the user + sees a column and concludes the metric is broken, not the call. Fail + loudly at the call site that can still be fixed instead. + """ + coerced = value.detach().cpu().numpy() if hasattr(value, 'detach') else normalize(value) + if coerced is None: + raise TypeError( + f"save_signals: signal {name!r} has unsupported type " + f"{type(value).__name__}. Pass a torch.Tensor, a numpy array, a " + f"list, or a scalar -- one value per sample, aligned with " + f"batch_ids. (A value of None would create the column and never " + f"write to it.)" + ) + return _reduce_signal(coerced) + # Processing signals if isinstance(signals, dict): losses_data = { - 'signals//' + k: (lambda arr: arr.mean(axis=tuple(range(1, arr.ndim))) if isinstance(arr, np.ndarray) and arr.ndim > 1 else arr)( - v.detach().cpu().numpy() if hasattr(v, 'detach') else normalize(v) - ) - for k, v in signals.items() + 'signals//' + k: _coerce_signal(k, v) for k, v in signals.items() } elif signals is not None and isinstance(signals, (th.Tensor, np.ndarray, list)): - losses_data = { - "signals//default": (lambda arr: arr.mean(axis=tuple(range(1, arr.ndim))) if isinstance(arr, np.ndarray) and arr.ndim > 1 else arr)( - signals.detach().cpu().numpy() if hasattr(signals, 'detach') else normalize(signals) - ) - } + losses_data = {"signals//default": _coerce_signal("default", signals)} else: losses_data = None - # Enqueue to dataframe manager buffer for efficiency + # Enqueue to dataframe manager buffer for efficiency. + # + # _seen=False (internal) writes the values WITHOUT a step, which is what + # bumps a sample's nb_seen/last_seen: projection coordinates describe where + # a sample sits in a picture, not a step at which the model saw it. DATAFRAME_M.enqueue_batch( sample_ids=batch_ids_np, preds_raw=preds_raw_np if preds_raw_np is not None else preds_raw, preds=preds_np if preds_np is not None else preds, targets=target_np if target_np is not None else targets, losses=losses_data, - step=step + step=step if _seen else None ) + # Live 3-D projection, for the usecases that never register a flag="loss" + # criterion and write their per-sample values straight here (YOLO/Ultralytics + # wrappers, custom eval loops, anything hand-rolled). Between this and the + # loss wrapper, every supported way of producing a per-sample value feeds + # the projection. Re-entrancy safe: the projection's own coordinate + # write-back lands in this same function and is guarded there. Not for a + # _seen=False write, which carries no batch the model just ran. + if _seen: + _projection_observe_batch(batch_ids, step) + # Reactive signals: these just-logged signals may satisfy an inputs=[...] # signal. Only logged (queryable) signals can be inputs. _react=False on the # reactive persist path prevents re-entrant dispatch. @@ -3315,7 +3441,20 @@ def track_model_signals(model=None, **kwargs): def clear_all(): """Clear all WeightsLab registries (models, dataloaders, etc.).""" + global DATAFRAME_M ledgers.clear_all() + # The cached dataframe handle is a ledger Proxy that clear_all just + # orphaned: the next registration creates a NEW proxy, and save_signals + # kept writing to the dead one ("Proxy target not set"). Re-resolve lazily. + DATAFRAME_M = None + # The projection hooks the model the ledger just let go of; a later wrap + # would re-attach anyway, but until then nothing should keep capturing its + # forward passes, nor apply a parked checkpoint restore to the next model. + _detach_projection() + from weightslab.projection import clear_pending_restore + from weightslab.projection.registry import clear_registry + clear_pending_restore() + clear_registry() def _unpack_batch(batch, device=None): @@ -3377,6 +3516,61 @@ def _unpack_batch(batch, device=None): # EVALUATION MODE PUBLIC API # ############################################################################################################## +def _resolve_module_device(module): + """Where a module's tensors live: parameters, else buffers, else the GPU. + + Returns None only when there is genuinely nothing to go on (no torch, no + tensors, no accelerator) — in which case leaving the batch where it is, is + the only honest option. + """ + try: + import torch as _th + except Exception: + return None + + for getter in ("parameters", "buffers"): + try: + fn = getattr(module, getter, None) + if fn is None: + continue + for tensor in fn(): + return tensor.device + except Exception: + continue + + # A module with neither (or a proxy that forwards neither) still has to + # meet its batch somewhere: prefer the accelerator, since a CPU batch + # reaching a CUDA model is the failure this exists to prevent. + try: + if _th.cuda.is_available(): + return _th.device("cuda") + except Exception: + pass + return None + + +def _move_to_device(value, device): + """Move tensors to ``device``, descending into tuples/lists/dicts. + + Anything that is not a tensor (ids, strings, metadata dicts) is returned + unchanged, and a failure to move is not fatal — the caller still reports + the forward error, which is more informative than a half-moved batch. + """ + if device is None or value is None: + return value + if hasattr(value, "to") and hasattr(value, "device"): + try: + return value.to(device) + except Exception: + return value + if isinstance(value, dict): + return {k: _move_to_device(v, device) for k, v in value.items()} + if isinstance(value, (list, tuple)): + moved = [_move_to_device(v, device) for v in value] + return type(value)(moved) if isinstance(value, list) else tuple(moved) + return value + + def _make_default_eval_fn(model): """Return a default evaluation callable that uses all registered ledger signals. @@ -3411,11 +3605,28 @@ def _default_eval(loader): no_grad_ctx = None try: - device = None - try: - device = next(model.parameters()).device - except (StopIteration, Exception): - pass + # Resolve where the model actually lives, and put the batch there. + # + # This used to read the first parameter and swallow every failure, + # leaving device=None — and a None device means the batch is never + # moved, so a CUDA model meets a CPU batch and every forward dies + # with "found at least two devices". Fall back through buffers (a + # model can be all-buffers) and then to the accelerator, so the + # answer is CUDA whenever CUDA is in play. + device = _resolve_module_device(model) + + # Re-home the module onto that device. Non-persistent buffers are + # absent from a state_dict, so a checkpoint restore can leave them + # behind on CPU while the parameters are on CUDA; .to() is a no-op + # for anything already there and fixes exactly that case. + if device is not None: + try: + model.to(device) + except Exception: + pass + + # Signals retired mid-pass because they raised (see below). + _failed_signals = set() # Resolve registered signals once before the loop. signal_names = [] @@ -3431,22 +3642,19 @@ def _default_eval(loader): if inputs is None: continue - if device is not None and hasattr(inputs, "to"): - try: - inputs = inputs.to(device) - except Exception: - pass - if targets is not None and device is not None and hasattr(targets, "to"): - try: - targets = targets.to(device) - except Exception: - pass + # Move recursively: a model can take a tuple/list/dict of + # tensors, and a top-level hasattr(x, "to") check silently + # leaves every one of those on the CPU. + inputs = _move_to_device(inputs, device) + targets = _move_to_device(targets, device) preds = model(inputs) # infer predictions # Call each registered signal so its wrapped forward/compute # fires and feeds into the evaluation-mode logger buffer. for sig_name in signal_names: + if sig_name in _failed_signals: + continue try: sig = get_signal(sig_name) if sig is None: @@ -3480,12 +3688,43 @@ def _default_eval(loader): if not attempted: sig(preds) except Exception as _se: - logger.debug( - "[wl.default_eval] signal '%s' failed: %s\nAre you sure signal {%s} is compatible with weightslab?", sig_name, _se, sig_name + # Retire the signal for the rest of this pass. + # + # These calls hand a signal the model's RAW output, + # which not every metric accepts: a binary metric + # given [B, num_classes] indexes out of bounds. On + # CPU that is the IndexError we just caught; on + # CUDA the same index is a device-side assert that + # poisons the context for the WHOLE process, so the + # next training step dies somewhere unrelated. + # Repeating it once per batch turns one bad + # contract into a guaranteed kill. + _failed_signals.add(sig_name) + logger.warning( + "[wl.default_eval] signal '%s' failed and is " + "skipped for this evaluation: %s | preds=%s " + "targets=%s. Signals are called with the " + "model's raw output — a metric that needs a " + "different shape must adapt it in update().", + sig_name, _se, + tuple(getattr(preds, "shape", ())) or "n/a", + tuple(getattr(targets, "shape", ())) or "n/a", ) except Exception as _be: - logger.debug("[wl.default_eval] batch forward failed: %s", _be) + # A device mismatch here used to surface as a bare debug + # line repeated once per batch, with nothing saying WHICH + # tensor was where. Name the devices so the next one is + # diagnosable from the log alone. + _where = "" + if "two devices" in str(_be): + try: + _where = (f" (model on {_resolve_module_device(model)}, " + f"inputs on {getattr(inputs, 'device', 'n/a')})") + except Exception: + pass + logger.warning( + "[wl.default_eval] batch forward failed: %s%s", _be, _where) finally: if no_grad_ctx is not None: try: @@ -4211,12 +4450,20 @@ def __iter__(self): # that stuck flag makes the next plain `next(train_loader)` call # in the user's training loop wrongly propagate StopIteration at # the next epoch boundary instead of transparently auto-resetting. + # + # Release only THIS loop's hold (the proxy iterator's close()), never + # the flag outright: the training script may be mid-way through its + # own `for batch in` the same loader, and clearing the flag under it + # turned that loop endless. + close = getattr(it, 'close', None) underlying = getattr(it, '_it', None) - if underlying is not None: - try: + try: + if callable(close): + close() + elif underlying is not None: underlying.is_a_loop = False - except Exception: - pass + except Exception: + pass # ############################################################################################################## diff --git a/weightslab/trainer/services/agent_service.py b/weightslab/trainer/services/agent_service.py index 3403b370..4e0a0b66 100644 --- a/weightslab/trainer/services/agent_service.py +++ b/weightslab/trainer/services/agent_service.py @@ -61,9 +61,16 @@ def CheckAgentHealth(self, request, context): Check whether any LLM provider is ready to serve queries. Returns: - AgentHealthResponse { available: bool, message: str } + AgentHealthResponse { available: bool, message: str, + current_model: str } - available=True → "Ready to help you." - available=False → "Agent not configured. Type /init to set up." + + `current_model` is re-resolved from OpenCode's config rather than read + off a snapshot, so a model chosen from the CLI (or by another studio + tab) is visible here. Health is polled, which makes this the channel + by which the picker learns it is out of date — before, a client could + only know the model IT had set, so a CLI switch never reached the UI. """ available = self._is_available() message = ( @@ -72,7 +79,9 @@ def CheckAgentHealth(self, request, context): else "Agent not configured. Type /init to set up." ) logger.debug("CheckAgentHealth → available=%s", available) - return pb2.AgentHealthResponse(available=available, message=message) + return pb2.AgentHealthResponse( + available=available, message=message, + current_model=self._current_model()) @safe_grpc(lambda msg: pb2.InitializeAgentResponse(success=False, message=msg)) def InitializeAgent(self, request, context): @@ -114,6 +123,21 @@ def InitializeAgent(self, request, context): success, message = agent.initialize_with_cloud_key(request.api_key, provider_name, request.model) return pb2.InitializeAgentResponse(success=success, message=message) + def _current_model(self) -> str: + """The model the next query will use, or "" when it cannot be resolved. + + Never raises: this feeds a health poll, and an agent that is down must + report an empty model rather than fail the whole response. + """ + agent = self._agent + if agent is None: + return "" + try: + return str(agent.current_model() or "") + except Exception as exc: # noqa: BLE001 + logger.debug("[AgentService] current_model unavailable: %s", exc) + return "" + @safe_grpc(lambda msg: pb2.ChangeAgentModelResponse(success=False, message=msg)) def ChangeAgentModel(self, request, context): """ @@ -150,7 +174,9 @@ def GetAgentModels(self, request, context): ) ok, models, message = agent.get_available_models() - return pb2.GetAgentModelsResponse(success=ok, models=models, message=message) + return pb2.GetAgentModelsResponse( + success=ok, models=models, message=message, + current_model=self._current_model()) @safe_grpc(lambda msg: pb2.ResetAgentResponse(success=False, message=msg)) def ResetAgent(self, request, context): diff --git a/weightslab/trainer/services/data_image_utils.py b/weightslab/trainer/services/data_image_utils.py index 30695b83..a50bc366 100644 --- a/weightslab/trainer/services/data_image_utils.py +++ b/weightslab/trainer/services/data_image_utils.py @@ -197,7 +197,7 @@ def render_tabular_heatmap(vec: np.ndarray, cell: int = 16, quality: int = 60): norm = (grid - mn) / (mx - mn + 1e-8) img8 = np.clip(norm * 255.0, 0, 255).astype(np.uint8) try: - pil = Image.fromarray(img8, mode="L").resize( + pil = Image.fromarray(img8).resize( (side * cell, side * cell), Image.Resampling.NEAREST) buf = io.BytesIO() pil.convert("RGB").save(buf, format="WEBP", quality=quality, method=1) diff --git a/weightslab/trainer/services/data_service.py b/weightslab/trainer/services/data_service.py index 752de53a..19455299 100755 --- a/weightslab/trainer/services/data_service.py +++ b/weightslab/trainer/services/data_service.py @@ -41,7 +41,7 @@ is_point_cloud_detection_task, ) from weightslab.data.video_utils import ( - describe_clip, is_generation_task, is_video_task, + describe_clip, has_playable_media, is_generation_task, select_frame_annotation, ) from weightslab.data import media_store from weightslab.trainer.trainer_tools import execute_df_operation, generate_overview, encode_image_to_raw_bytes @@ -135,6 +135,21 @@ def _point_cloud_chunk_bytes() -> int: _DEFAULT_MEDIA_CHUNK_BYTES = 1 << 18 # 256 KiB +def _media_entry_bytes(value) -> int: + """Payload size of a cached media entry, for the media cache byte budget. + + Entries are dicts whose "data" holds the encoded container; ``len()`` on + the dict itself would count keys, not bytes. Module-level rather than a + method so anything reusing ``_media_cache_put`` (the gRPC test stubs bind + it onto their own class) keeps working. + """ + payload = value.get("data") if isinstance(value, dict) else value + try: + return len(payload) + except TypeError: + return 0 + + def _media_chunk_bytes() -> int: """Read WL_MEDIA_CHUNK_BYTES; non-positive/invalid falls back to the default.""" return _env_int("WL_MEDIA_CHUNK_BYTES", _DEFAULT_MEDIA_CHUNK_BYTES) @@ -147,6 +162,124 @@ def _media_chunk_bytes() -> int: # text/scalar branch alongside classification and tabular. _NON_MASK_TASKS = ("classification", "tabular") +# Guards the lazy creation of a DataService's projection cache, so two gRPC +# workers asking at once do not each build the first index. +_PROJECTION_CACHE_LOCK = threading.Lock() + + +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. + + Vectorised for the dtypes that actually turn up, and only those; anything + else still goes through ``is_set_flag`` one value at a time, so the rules + that function documents keep holding exactly. + + Worth the branching: this is called once per point on every projection + refresh, and the per-element path was the single largest cost of that RPC + -- about 15 ms of a 34 ms response over 17k points, repeated every time the + camera settles and every ten seconds besides. + """ + if series is None: + return np.zeros(0, dtype=bool) + + dtype = getattr(series, "dtype", None) + + # pandas nullable boolean: NA is not set, everything else is itself. + if isinstance(dtype, pd.BooleanDtype): + return series.fillna(False).to_numpy(dtype=bool) + + if dtype is not None and not isinstance(dtype, pd.CategoricalDtype): + kind = getattr(dtype, "kind", None) + # Plain numpy bool: no missing value is representable. + if kind == "b": + return series.to_numpy(dtype=bool) + # Numeric: NaN is missing (NOT True, which is what bool() would say), + # and zero is false. + if kind in ("i", "u"): + return series.to_numpy() != 0 + if kind == "f": + values = series.to_numpy(dtype=np.float64) + return np.isfinite(values) & (values != 0) + # pandas nullable Int64/Float64 and friends. + if isinstance(dtype, pd.api.extensions.ExtensionDtype) and \ + pd.api.types.is_numeric_dtype(dtype): + values = pd.to_numeric(series, errors="coerce").to_numpy(dtype=np.float64) + return np.isfinite(values) & (values != 0) + + return np.fromiter((is_set_flag(v) for v in series.tolist()), + dtype=bool, count=len(series)) + + +def _boxes_to_bbox_payload(boxes): + """``[N, >=4]`` rows -> the JSON payload the 2-D box renderer reads. + + Shared by the full record path and the preview-cache path, which used to + disagree: the preview path had no plain-detection branch at all, so a + box array fell through to the segmentation encoder and was PIL-resized + into a 64x64 RLE "mask". The first fetch of a sample looked right (it came + from the full path) and every fetch after it was served from the poisoned + cache. + """ + arr = to_numpy_safe(boxes) + if arr is None: + try: + arr = np.asarray(boxes, dtype=np.float32) + except Exception: + return None + if arr.ndim == 1 and arr.size >= 4: + arr = arr.reshape(1, -1) + if arr.size == 0 or arr.ndim != 2 or arr.shape[-1] < 4: + return None + return {"bboxes": arr.tolist(), "format": detect_bbox_format(arr[..., :4])} + + +def _video_frame_count(dataset, task_type=None) -> int: + """Frames in this dataset's clips, or 0 when it is not a video dataset. + + Read from the dataset's declared attributes (never by decoding), because + this runs per rendered row. 0 means "do not treat annotations as + per-frame", which is the right answer for every image dataset. + """ + if dataset is None or not has_playable_media(dataset, task_type): + return 0 + try: + return int(describe_clip(dataset).get("frame_count") or 0) + except Exception: + return 0 + def is_set_flag(value) -> bool: """True only for a boolean column value that is really SET and true. @@ -832,7 +965,21 @@ def _build_preview_cache(self) -> None: value_string=json.dumps(payload))) except Exception as exc: logger.debug("[PreviewCache] detection_pointcloud target skipped: %s", exc) + elif label is not None and _early_task == "detection": + payload = _boxes_to_bbox_payload(select_frame_annotation( + label, _video_frame_count(dataset, _early_task))) + if payload: + stats.append(create_data_stat( + 'target', 'string', shape=[1], + value_string=json.dumps(payload))) elif label is not None: + # Video segmentation annotates every frame. Reduce to + # the poster frame FIRST: the ndim>2 branch below reads + # a 3-D array as [H, W, C] and takes [:, :, 0], which + # on a [T, H, W] mask stack is a frames x height strip + # — a silently wrong overlay, not an error. + label = select_frame_annotation( + label, _video_frame_count(dataset, _early_task)) label_arr = to_numpy_safe(label) if label_arr is None: try: @@ -864,7 +1011,17 @@ def _build_preview_cache(self) -> None: value_string=json.dumps(payload))) except Exception as exc: logger.debug("[PreviewCache] detection_pointcloud pred skipped: %s", exc) + elif pred is not None and _early_task == "detection": + payload = _boxes_to_bbox_payload(select_frame_annotation( + pred, _video_frame_count(dataset, _early_task))) + if payload: + stats.append(create_data_stat( + 'pred', 'string', shape=[1], + value_string=json.dumps(payload))) elif pred is not None: + # Same poster-frame reduction as the GT above. + pred = select_frame_annotation( + pred, _video_frame_count(dataset, _early_task)) pred_arr = to_numpy_safe(pred) if pred_arr is None: try: @@ -946,6 +1103,13 @@ def _refresh_preview_masks_from_row(self, preview_rec: "pb2.DataRecord", row: pd _task = next((st.value_string for st in rec.data_stats if st.name == "task_type"), "") if is_point_cloud_detection_task(_task): return self._refresh_preview_boxes_3d_from_row(rec, row) + # Plain 2-D detection carries box rows, which are no more rasterizable + # than point-cloud boxes are. Without this they fell through to the + # mask encoder below and a [N, 6] box array was PIL-resized into a + # 64x64 RLE "mask" — so a detection sample rendered correctly on its + # first fetch and as garbage from the preview cache thereafter. + if str(_task).strip().lower() == "detection": + return self._refresh_preview_boxes_2d_from_row(rec, row) # Infer preview dimensions from cached raw_data shape; fallback to 64x64. target_h, target_w = 64, 64 @@ -998,6 +1162,41 @@ def _upsert_mask_stat(stat_name: str, rle_bytes: bytes | None, shape: list[int] return rec + def _refresh_preview_boxes_2d_from_row(self, rec: "pb2.DataRecord", row: pd.Series) -> "pb2.DataRecord": + """Refresh GT/pred box JSON of a plain (2-D) detection preview record. + + Video detection annotates every frame, so the rows are reduced to the + poster frame first — the same frame the cached thumbnail shows. + """ + origin = next((st.value_string for st in rec.data_stats if st.name == "origin"), None) + dataset = self._get_dataset(origin) if origin else None + frames = _video_frame_count(dataset, "detection") + + def _upsert_box_stat(stat_name: str, value_like) -> None: + payload = _boxes_to_bbox_payload( + select_frame_annotation(value_like, frames)) + if not payload: + return + new_stat = create_data_stat( + stat_name, 'string', shape=[1], value_string=json.dumps(payload)) + for i, st in enumerate(rec.data_stats): + if st.name == stat_name: + rec.data_stats[i].CopyFrom(new_stat) + return + rec.data_stats.append(new_stat) + + # A stale mask stat from an earlier encode would otherwise sit next to + # the boxes and be drawn as an overlay. + for name in ("target", "pred_mask"): + for i, st in enumerate(rec.data_stats): + if st.name == name and st.type == "rle_mask": + del rec.data_stats[i] + break + + _upsert_box_stat('target', row.get(SampleStatsEx.TARGET.value)) + _upsert_box_stat('pred', row.get(SampleStatsEx.PREDICTION.value)) + return rec + def _refresh_preview_boxes_3d_from_row(self, rec: "pb2.DataRecord", row: pd.Series) -> "pb2.DataRecord": """Refresh GT/pred box JSON stats of a detection_pointcloud preview record.""" from weightslab.data.point_cloud_utils import serialize_pointcloud_box_payload @@ -1674,6 +1873,13 @@ def _process_sample_row(self, args): if label_raw is None and dataset is not None: label_raw = load_label(dataset, sample_id) + # Video segmentation/detection annotates every frame, but the + # grid and the modal's still view draw exactly one — the + # poster. Take that frame's annotation, or the renderer gets a + # 3-D stack RLE'd into a stat it reads as 2-D. + label_raw = select_frame_annotation( + label_raw, _video_frame_count(dataset, task_type)) + # Handle 3D (point cloud) detection: metric boxes are sent both # as raw 3D rows (for the interactive viewer) and projected to # the BEV image frame (legacy 'bboxes' key, so the existing 2D @@ -1944,6 +2150,11 @@ def _json_default(o): if not _is_non_mask_task(task_type) and pred is not None: t0_pmask = time.time() + # Same poster-frame rule as the ground truth above: a + # prediction drawn on the still must describe the frame the + # still actually shows. + pred = select_frame_annotation( + pred, _video_frame_count(dataset, task_type)) # 3D (point cloud) detection predictions: same dual payload as # the GT path (BEV-projected 'bboxes' + raw 'bboxes_3d'). @@ -2322,7 +2533,14 @@ def _json_default(o): # frame, so advertise the clip's shape here. This lets the # grid badge the cell and the modal size its player without # anyone paying for a GetMedia round-trip first. - if is_video_task(task_type): + # Keyed on the dataset's media kind, not just the task + # type: a video *classification* dataset declares + # media_kind="video" and keeps its own task type, and its + # poster already routes through is_video_sample. Gating + # this on video_generation alone left those samples + # playable over GetMedia but never advertised as such, so + # the UI rendered them as stills. + if has_playable_media(dataset, task_type): media_info = describe_clip(dataset) if media_info: data_stats.append( @@ -2441,6 +2659,42 @@ def _get_categorical_tag_defs(self) -> List["pb2.CategoricalTagDef"]: logger.debug(f"Error building categorical tag defs: {e}") return defs + def _apply_projection_selection(self, query: str) -> pb2.DataQueryResponse: + """Narrow the data board to a lasso selection made in the projection + (``@projection_selection ``). + + The selection is computed and kept by the projection cache because it + can hold millions of samples -- the drawn ones and the many the + level-of-detail view never sent -- which no query string can carry. A + ``sample_id.isin([...])`` of a million ids took longer to parse than + the rest of the round trip. + """ + parts = query.split() + token = parts[1] if len(parts) > 1 else "" + cache = getattr(self, "_projection_cache", None) + if cache is None or not token or cache.selection(token) is None: + return pb2.DataQueryResponse( + success=False, + message="That projection selection has expired -- draw the lasso again.") + # Picked from the whole dataset: the live view when nothing narrows it, + # else the frame the projection indexed the whole dataset from. Either + # is usually the very frame the lasso was drawn on, whose rows the + # selection already holds. + base = self._all_datasets_df + if self._is_filtered or base is None or base.empty: + base = cache.full_frame() + if base is None: + base = self._pull_into_all_data_view_df() + df = base.iloc[cache.selection_rows(token, base)] + with self._watched_lock("_lock[ApplyDataQuery/projection]"): + self._all_datasets_df = df + self._is_filtered = True + return self._build_success_response( + df=df, + message=f"{len(df)} samples selected in the projection", + intent_type=pb2.INTENT_FILTER, + ) + def _build_success_response( self, df, @@ -4139,32 +4393,54 @@ def _fastUpdateInternals(self, max_dirty: int = 250_000) -> bool: # 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)) + # + # The view's keys as strings are built once per view INDEX, not once + # per call: converting every row cost ~3 s per refresh at 10M samples, + # for a delta of a few thousand. An Index is immutable and every new + # view brings a new one, so identity is a safe cache key. + _cached = getattr(self, "_fast_view_keys", None) + if _cached is not None and _cached[0] is view.index: + _view_keys = _cached[1] + else: + 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) + _view_keys = pd.Index(view_keys.astype(str)) + self._fast_view_keys = (view.index, _view_keys) _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] + if _view_keys.is_unique: + # O(delta): each dirty sample looked up in the view's own hash + # index (built once with the cached keys above, then reused). + _pos = _view_keys.get_indexer(_sub_keys) + _found = _pos >= 0 + if not _found.any(): + 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 not _found.all(): + return False + _rows = _pos + _take = np.arange(len(_pos)) + else: + # 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. + _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 + 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) _vals = sub[c].to_numpy()[_take] @@ -5255,6 +5531,10 @@ def _run_ops(target): success=True, message="View has been reset successfully.", ) + elif request.query.strip().lower().startswith("@projection_selection"): + logger.info(f"[ApplyDataQuery] BYPASSING AGENT - Projection selection: {request.query[:100]}") + return self._apply_projection_selection(request.query) + elif request.query.lower().replace("'''", "\"\"\"").replace('\"\"\"', "").replace('\'\'\'', "").replace(" ", "").startswith("@overview"): logger.info(f"[ApplyDataQuery] BYPASSING AGENT - Direct overview operation: {request.query[:100]}...") # Force view reset @@ -5421,6 +5701,60 @@ def GetDataSamples(self, request, context): data_records=[] ) + def GetProjection(self, request, context): + """Points of the live 3-D projection inside the client's view box. + + Follows the data board's filter when ``request.follow_view`` is set, + which is the ordinary case: filtering the view (by hand or via the + agent) narrows the projection with it. + + The exception is a filter the projection itself produced. A lasso + pushes a sample_id query into the data board; re-reading that here + would collapse the cloud to the points just selected and leave nothing + to select from next, so the client clears the flag for that one + refresh and the whole dataset is served instead. + + Grid Overview (``follow_view`` with ``restrict_sample_ids``) is the + exception the other way: the cloud is exactly the samples the grid is + rendering, served whole, whether or not the data view is also filtered + -- what is on screen is the truth, and a subset's page is a subset of + that subset. + + All the slicing/decimation lives in projection_service; this is the RPC + shell around it. The index it answers from is built and refreshed by + ``ProjectionCache``, outside ``_lock``: holding the lock for a whole + projection request stalled every other board for as long as it took. + """ + try: + from weightslab.trainer.services import projection_service + if self._all_datasets_df is None: + self._initialize_data_service() + cache = getattr(self, "_projection_cache", None) + if cache is None: + with _PROJECTION_CACHE_LOCK: + cache = getattr(self, "_projection_cache", None) + if cache is None: + cache = self._projection_cache = projection_service.ProjectionCache() + + def refreshed_view(): + # The view brought up to date first: what an index is built + # from. Called when the cache (re)builds, not on every request. + with self._watched_lock("_lock[GetProjection]"): + if not self._fastUpdateInternals(): + self._slowUpdateInternals() + return self._all_datasets_df + + return cache.serve( + request, + view=lambda: self._all_datasets_df, + is_filtered=lambda: bool(self._is_filtered), + refresh_view=refreshed_view, + pull_full=self._pull_into_all_data_view_df, + ) + except Exception as e: + logger.error("Error in GetProjection: %s", str(e), exc_info=True) + return pb2.ProjectionResponse(success=False, message=str(e)) + def GetHistogram(self, request, context): """Server-side histogram binning of one column (typed RPC). @@ -5448,10 +5782,10 @@ 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) + # 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 @@ -5602,9 +5936,16 @@ def _field(frame, name): # Same, for GetMedia (already-compressed MP4/WAV bytes per message). _MEDIA_CHUNK_BYTES = _media_chunk_bytes() - # How many encoded clips to keep. Clips are megabytes each, so this only - # needs to cover "the sample the user is currently looking at". - _MEDIA_CACHE_ENTRIES = 3 + # How many encoded clips to keep, and how many bytes they may occupy. + # + # This used to be 3, sized for "the sample the user is looking at". That + # made the studio's modal prefetch window (6 neighbours) evict its own + # entries before they could be used, and re-muxing a clip costs an ffmpeg + # round-trip on the training process. Sized now to cover a browsing window; + # the byte budget is what actually bounds memory, since clip sizes vary by + # orders of magnitude between datasets. + _MEDIA_CACHE_ENTRIES = _env_int("WL_MEDIA_CACHE_ENTRIES", 32) + _MEDIA_CACHE_BYTES = _env_int("WL_MEDIA_CACHE_MB", 512) * 1024 * 1024 def _locate_sample(self, sample_id: int, origin: str = ""): """Resolve ``sample_id`` to ``(dataset, ds_index, member_rank)``. @@ -5789,15 +6130,23 @@ def _media_cache_get(self, key): def _media_cache_put(self, key, value) -> None: """Insert into the encoded-media LRU, evicting the oldest entries. - Encoded clips are large, so this is deliberately tiny — it exists to - make re-opening and scrubbing the *same* sample free, not to hold a - working set. + Bounded by BOTH entry count and total bytes: the count keeps a + browsing window resident so the studio's prefetch is not wasted, and + the byte budget is what stops a dataset of large clips from pinning + the training process's memory. The newest entry is never evicted — + someone is about to stream it. """ with self._media_cache_lock: self._media_cache.pop(key, None) # keep re-puts from double-counting self._media_cache[key] = value - while len(self._media_cache) > self._MEDIA_CACHE_ENTRIES: - self._media_cache.pop(next(iter(self._media_cache))) + total = sum(_media_entry_bytes(v) for v in self._media_cache.values()) + while len(self._media_cache) > 1 and ( + len(self._media_cache) > self._MEDIA_CACHE_ENTRIES + or total > self._MEDIA_CACHE_BYTES): + oldest = next(iter(self._media_cache)) + if oldest == key: + break + total -= _media_entry_bytes(self._media_cache.pop(oldest)) def GetPointCloud(self, request, context): """Stream one sample's raw point cloud as binary float32 chunks. @@ -6026,6 +6375,24 @@ def _register_tag(self, tag_name: str): ) def EditDataSample(self, request, context): + """Edit sample metadata; see :meth:`_edit_data_sample`. + + A successful edit also tells the projection, whose whole-dataset cloud + may be drawn from a copy pulled before it (ProjectionCache.note_edit): + a sample discarded in the grid is greyed in the projection within + seconds, not at the next scheduled pull. + """ + response = self._edit_data_sample(request, context) + if getattr(response, "success", False): + cache = getattr(self, "_projection_cache", None) + if cache is not None: + try: + cache.note_edit() + except Exception as exc: + logger.debug(f"[projection] could not mark the cache stale: {exc}") + return response + + def _edit_data_sample(self, request, context): """ Edit sample metadata (tags and discarded). diff --git a/weightslab/trainer/services/notebook_service.py b/weightslab/trainer/services/notebook_service.py index 9a3c24e5..fed952ad 100644 --- a/weightslab/trainer/services/notebook_service.py +++ b/weightslab/trainer/services/notebook_service.py @@ -25,6 +25,7 @@ """ import io +import gc import os import re import ast @@ -37,8 +38,10 @@ import shutil import logging import builtins +import hashlib import threading import traceback +from concurrent.futures import Future, TimeoutError as FuturesTimeoutError import contextlib from pathlib import Path from concurrent.futures import ThreadPoolExecutor @@ -51,6 +54,17 @@ logger = logging.getLogger(__name__) +# How long a finished GenerateNotebookCode result stays collectable by a client +# that reconnects (a browser refresh mid-generation). Long enough to cover a +# reload plus the user reopening the notebook, short enough that a later +# identical re-ask is a genuine new generation rather than a replay. +_GEN_CACHE_TTL_S = float(os.environ.get("WL_NOTEBOOK_GEN_TTL_S", "900")) +# Longest a single RPC blocks waiting on a generation. The OpenCode chain has +# its own (shorter) timeout, so reaching this means something is wedged; the +# entry stays cached either way, so the client can re-ask and collect it. +_GEN_WAIT_TIMEOUT_S = float(os.environ.get("WL_NOTEBOOK_GEN_WAIT_S", "180")) + + # The notebook file lives directly under root_log_dir so it travels with the # experiment's checkpoints/logs. NOTEBOOK_FILENAME = "notebook.ipynb" @@ -333,6 +347,15 @@ def _shutdown_embedded_kernel(timeout: float = 5.0) -> None: atexit runs before interpreter finalization, so stopping the loop here lets app.start() return and the thread exit while logging still works. + + It must also run BEFORE ipykernel's own atexit hooks: ``app.initialize()`` + registers ``app.close`` (closes every socket and terminates the zmq + context) and ``IOPubThread.stop``. atexit runs handlers last-registered + first, so _run_embedded_kernel registers this one right after + initialize(). Registered before it (it used to be, at thread creation), + app.close ran first and pulled the sockets out from under the + still-running loop -- the same warnings as above, then ``terminate called + without an active exception`` and a core dump on Linux CI. """ thread = _EMBED_STATE.get("thread") if thread is None or not thread.is_alive(): @@ -345,9 +368,11 @@ def _shutdown_embedded_kernel(timeout: float = 5.0) -> None: # -- measured: stored loop id ...692880 vs app's ...181904, and the thread # stayed alive through a 10s join. IOLoop.add_callback is documented # thread-safe, which is what lets us reach into it from here. + app = None try: from ipykernel.kernelapp import IPKernelApp - io_loop = getattr(IPKernelApp.instance(), "io_loop", None) + app = IPKernelApp.instance() + io_loop = getattr(app, "io_loop", None) if io_loop is not None: io_loop.add_callback(io_loop.stop) except Exception: @@ -358,6 +383,15 @@ def _shutdown_embedded_kernel(timeout: float = 5.0) -> None: # way out. Worst case we are back to the old behaviour. thread.join(timeout=timeout) + # IPythonKernel adds a gc callback and never removes it. A collection + # during finalization then calls it with the threading module already torn + # down: "Exception ignored while calling GC callback ... 'NoneType' object + # has no attribute 'enumerate'" after an otherwise clean exit. + try: + gc.callbacks.remove(app.kernel._clean_thread_parent_frames) + except (AttributeError, ValueError): + pass # no kernel, or an ipykernel that never registered it + def configure_embedded_kernel(enabled: bool) -> None: """Record wl.serve()'s enable/disable decision for the embedded kernel.""" @@ -428,11 +462,8 @@ def ensure_embedded_kernel(data_service, root_log_dir: Path) -> None: name="WL-Embedded-Jupyter-Kernel", daemon=True, ) _EMBED_STATE["thread"] = _kernel_thread - # Registered only once the kernel actually starts, and only here -- - # this is the one place that creates the thread. See - # _shutdown_embedded_kernel() for why a daemon thread left running into - # interpreter finalization aborts the process. - atexit.register(_shutdown_embedded_kernel) + # The atexit hook that stops this thread is registered by the thread + # itself, after ipykernel's own -- see _shutdown_embedded_kernel(). _kernel_thread.start() try: if _wait_for_connection_file(connection_file, timeout=15.0): @@ -591,6 +622,11 @@ def _run_embedded_kernel(connection_file: Path) -> None: app.init_signal = lambda: None try: app.initialize([]) + # After initialize(), which registers ipykernel's own atexit hooks, so + # this one runs before them (atexit is LIFO). See + # _shutdown_embedded_kernel() for why the order decides between a clean + # exit and a core dump. + atexit.register(_shutdown_embedded_kernel) # traitlets configures logging while building any Application, and # logging.config.dictConfig CLOSES every handler in the process on its # way in (_clearExistingHandlers -> logging.shutdown) without detaching @@ -618,7 +654,8 @@ def _run_embedded_kernel(connection_file: Path) -> None: # training progress bar disappeared from the terminal once the notebook # kernel was enabled. _install_thread_routed_streams(console_stdout, console_stderr) - app.shell.colors = "NoColor" + # Lowercase: IPython 9 deprecates "NoColor"; IPython 8 matches caselessly. + app.shell.colors = "nocolor" try: app.shell.enable_matplotlib("inline") except Exception: @@ -663,6 +700,59 @@ def _run_embedded_kernel(connection_file: Path) -> None: logger.error("Embedded notebook kernel crashed during startup", exc_info=True) +# How long a cell waits for the training step in flight to finish once it has +# paused training (see pause_training_for_cell). +_PAUSE_BARRIER_TIMEOUT_S = 30.0 + +PAUSED_FOR_CELL_NOTE = ( + "[WeightsLab] Training paused to run this cell -- press Play in the header " + "to resume.\n") + + +def pause_training_for_cell(timeout_s: float = _PAUSE_BARRIER_TIMEOUT_S): + """Pause a running training before a notebook cell executes. + + A cell runs in the SAME process as the training loop, against the live + model, dataframe and checkpoints. Letting the loop keep stepping underneath + it means a cell that reads features, computes a t-SNE, or inspects weights + sees a model that changes mid-cell -- the numbers describe no single state + of the run. So executing a cell pauses training first, exactly like the + header's pause button (same controller, so the UI shows it paused), and it + STAYS paused: resuming is the user's call, via Play. + + pause() only stops the loop at its next step; the step in flight runs to + the end, holding the training lock. Taking that lock once (and releasing it + at once) is the barrier that waits it out -- the next step then blocks in + wait_if_paused() before it can take the lock again. + + Returns the note to print in the cell's output when it paused a running + training, else None (already paused, or no training at all). + """ + try: + from weightslab.components.global_monitoring import ( + pause_controller, weightslab_rlock) + except Exception: + return None + try: + if pause_controller.is_paused(): + return None + pause_controller.pause() + except Exception as exc: + logger.warning("notebook: could not pause training before the cell: %s", exc) + return None + note = PAUSED_FOR_CELL_NOTE + try: + if weightslab_rlock.acquire(timeout=timeout_s): + weightslab_rlock.release() + else: + note += ("[WeightsLab] (the training step in progress had not finished " + f"after {timeout_s:.0f}s; it completes alongside this cell.)\n") + except Exception: + pass + logger.info("notebook: training paused to run a cell") + return note + + def _install_kernel_hooks(shell) -> None: """Bracket every cell with the same write-guard + df-refresh the legacy kernel applies -- but as IPython events, which run ON the kernel's own @@ -674,6 +764,12 @@ def _pre_execute(): # This hook runs on the thread that executes the cell -- the one # _ThreadRoutedStream lets through to the notebook. _CELL_THREAD["ident"] = threading.get_ident() + # Code reaching the embedded kernel some other way than the Studio's + # RunNotebookCell (`jupyter console --existing`) pauses training too. + # From the Studio this is a no-op: RunNotebookCell already paused it. + note = pause_training_for_cell() + if note: + print(note, end="") try: shell.user_ns["df"] = get_df(_ACTIVE_BINDING["data_service"]) except Exception: @@ -1051,6 +1147,53 @@ def _collect_figures(self): # Service # --------------------------------------------------------------------------- +# The default notebook's second cell: the model's layers, which one the live +# projection reads, and how to point it at another. Inert until LAYER is set: +# re-projecting reads a whole split, which no "run all" should start by itself. +_PROJECTION_LAYER_CELL = '''\ +# Projection: which layer the parametric UMAP reads, and how to change it. +# `model` is the live model's ledger handle; resolve_module finds the real +# nn.Module behind it, whose named_modules() are the layers to choose from. +from weightslab.backend import ledgers + +net = wl.projection.resolve_module(model) if model is not None else None +live = wl.projection.get_tracker() +current = getattr(live, "layer_name", None) +if live is not None: + print(f"The projection reads {current} ({getattr(live, 'layer_detail', '')})" + f" -> '{live.signal_prefix}'\\n") +else: + print("No live projection in this run (projection=False, or no model yet).\\n") + +# a) Every layer, by index and name. +layers = [(n, m) for n, m in net.named_modules() if n] if net is not None else [] +for i, (name, module) in enumerate(layers): + own = sum(p.numel() for p in module.parameters(recurse=False)) + mark = " <- projection" if name == current else "" + print(f"{i:4d} {name:<40} {type(module).__name__:<18} {own:>12,} params{mark}") + +# b) Re-project from another layer: set LAYER to a name above, or to its +# index, and run this cell again. It reads that layer's features over the +# whole split and fits the UMAP encoder on them. Written to the live +# prefix, training then carries on from this fit (on a large dataset, pass +# max_samples= to fit on fewer samples first). +LAYER = None # e.g. "fc1", or 12 +loaders = ledgers.list_dataloaders() +SPLIT = "train_loader" if "train_loader" in loaders else (loaders[0] if loaders else None) + +if LAYER is not None: + name = layers[LAYER][0] if isinstance(LAYER, int) else LAYER + result = wl.project_dataset( + model, ledgers.get_dataloader(SPLIT), layer=name, + prefix=live.signal_prefix if live is not None else "umap", + # prefix="umap_" + name.replace(".", "_"), # or keep both, to compare in the board's picker + ) + print(f"{result['samples']:,} {SPLIT} samples re-projected from {name} -> {result['columns']}") + # For this run. To start from that layer next time, give it to the script: + # wl.watch_or_edit(model, flag="model", projection={"layer": name}) +''' + + class NotebookService: """gRPC facade for notebook cell execution, persistence, and code generation.""" @@ -1062,6 +1205,10 @@ def __init__(self, data_service, root_log_dir: str = None): # Base name (file stem, no .ipynb) of the notebook currently in use this # session. None until the first Get/Save resolves it. self._active_name = None + # In-flight and recently-finished GenerateNotebookCode calls, keyed by + # the hash of (prompt, context_code). See GenerateNotebookCode. + self._gen_lock = threading.Lock() + self._generations = {} # Eager: start (or rebind) the embedded kernel as soon as we know # data_service/root_log_dir, so external tools (`jupyter console # --existing `) and the studio panel can both attach without @@ -1164,8 +1311,9 @@ def _audit(self, action_type, status, details=None, error=None): @staticmethod def _default_notebook() -> dict: - """A minimal nbformat v4 notebook shown on first open: one worked - example (df + logger + checkpoint hash, all live ledger calls) and one + """A minimal nbformat v4 notebook shown on first open: two worked + examples -- df + logger + checkpoint hash, then the model's layers and + which one the projection reads (see _PROJECTION_LAYER_CELL) -- and one blank cell whose placeholder text (rendered client-side, see the studio's CodeMirror setup) demonstrates the "> ..." agent convention.""" def code(src): @@ -1219,6 +1367,7 @@ def code(src): "# checkpoint manager (also a live ledger call, like df and logger above).\n" "print(\"Current experiment hash:\", cm.get_current_experiment_hash() if cm is not None else None)" ), + code(_PROJECTION_LAYER_CELL), code(""), ] return { @@ -1248,6 +1397,12 @@ def RunNotebookCell(self, request, context): started = time.perf_counter() exec_count = 0 ok = True + # Before the kernel starts: the cell must see one state of the run, not + # a model the training loop keeps stepping underneath it. + note = pause_training_for_cell() + if note: + self._audit("notebook_pause_training", "success", {"cell_id": cell_id}) + yield pb2.NotebookCellChunk(cell_id=cell_id, stdout=note) try: kernel = self._get_kernel() for kind, payload in kernel.run_streaming(code): @@ -1329,14 +1484,95 @@ def SaveNotebook(self, request, context): @safe_grpc(lambda msg: pb2.GenerateNotebookCodeResponse(code="", explanation="", ok=False, error=msg)) def GenerateNotebookCode(self, request, context): + """Propose code for a ">" notebook cell. + + Deduplicated and cached by the hash of (prompt, context_code), which is + what lets a generation survive a browser refresh. The call used to be a + plain blocking invoke: refreshing the studio cancelled the grpc-web + request, but gRPC never interrupts a running sync handler, so OpenCode + finished the work, the audit log recorded a success, and the generated + code was written to a stream nobody was reading -- unrecoverable, since + OpenCodeChat creates a fresh session per call and agent.history keeps + only breadcrumbs, never the code. + + Now the work runs on its own thread behind a Future, and the request + key is what identifies it. The reloaded page replays the identical + request, so it either attaches to the still-running generation or picks + up its finished result immediately. Entries expire after + _GEN_CACHE_TTL_S so a stale key can never shadow a genuine re-ask. + """ agent = self._agent if agent is None: return pb2.GenerateNotebookCodeResponse( code="", explanation="", ok=False, error="Agent backend is not running.") + + prompt = request.prompt or "" + context_code = request.context_code or "" + key = hashlib.sha256( + f"{len(prompt)}:{prompt}{context_code}".encode("utf-8", "replace")).hexdigest() + + future, started = self._generation_future(agent, key, prompt, context_code) + if started: + logger.debug("GenerateNotebookCode: started generation %s", key[:8]) + else: + logger.debug("GenerateNotebookCode: joined existing generation %s", key[:8]) + try: - code, explanation = agent.generate_code(request.prompt or "", request.context_code or "") + code, explanation = future.result(timeout=_GEN_WAIT_TIMEOUT_S) + except FuturesTimeoutError: + # The work is still running and stays cached -- the client can ask + # again with the same key and pick it up rather than restarting it. + return pb2.GenerateNotebookCodeResponse( + code="", explanation="", ok=False, + error="Still generating — this cell will pick up the result shortly.") except Exception as exc: logger.info("GenerateNotebookCode failed: %s", exc) - return pb2.GenerateNotebookCodeResponse(code="", explanation="", ok=False, error=str(exc)) - self._audit("notebook_generate_code", "success", {"prompt": (request.prompt or "")[:500]}) - return pb2.GenerateNotebookCodeResponse(code=code, explanation=explanation, ok=True, error="") + return pb2.GenerateNotebookCodeResponse( + code="", explanation="", ok=False, error=str(exc)) + + # Audited once per generation, by the thread that produced it, not once + # per client that collects it (see _run_generation). + return pb2.GenerateNotebookCodeResponse( + code=code, explanation=explanation, ok=True, error="") + + def _generation_future(self, agent, key, prompt, context_code): + """Future for this request key, starting the work if nobody has yet. + + Returns ``(future, started_here)``. Expired entries are swept on the + way through, which is enough: the cache only grows on real requests. + """ + now = time.time() + with self._gen_lock: + for stale in [k for k, e in self._generations.items() + if now - e["at"] > _GEN_CACHE_TTL_S and e["future"].done()]: + self._generations.pop(stale, None) + + entry = self._generations.get(key) + if entry is not None: + return entry["future"], False + + future = Future() + self._generations[key] = {"future": future, "at": now} + + thread = threading.Thread( + target=self._run_generation, + args=(agent, key, prompt, context_code, future), + name=f"wl-notebook-gen-{key[:8]}", + daemon=True, + ) + thread.start() + return future, True + + def _run_generation(self, agent, key, prompt, context_code, future): + """Run one generation to completion regardless of who is still waiting.""" + try: + code, explanation = agent.generate_code(prompt, context_code) + except Exception as exc: + # Drop the entry so a retry actually retries instead of replaying + # the same failure for the rest of the TTL. + with self._gen_lock: + self._generations.pop(key, None) + future.set_exception(exc) + return + self._audit("notebook_generate_code", "success", {"prompt": prompt[:500]}) + future.set_result((code, explanation)) diff --git a/weightslab/trainer/services/projection_service.py b/weightslab/trainer/services/projection_service.py new file mode 100644 index 00000000..a86a0c62 --- /dev/null +++ b/weightslab/trainer/services/projection_service.py @@ -0,0 +1,1563 @@ +"""Serving for the live 3-D projection (see :mod:`weightslab.projection`). + +The coordinates themselves are written during training as ordinary per-sample +signals (``signals//umap_x`` and friends), so this module is purely a read +path: it slices the current data view down to what one camera can actually +draw and hands back packed arrays. + +Three things here are worth knowing about: + +**Level of detail is server-side, driven by the client's view.** The client +never asks for "the dataset"; it asks for "what my camera can see, at most N +points" -- a box, plus the camera's frustum planes. Zooming in shrinks that +region, so the same N buys finer detail. A projection over millions of samples +therefore never crosses the wire in full, and the renderer's budget is a +constant. + +**The subsample is stable, not random.** Every point has one fixed place in a +priority order (a deterministic hash of its sample id, stratified by cluster), +and a request returns the first N points of that order inside its region. So +the same request returns the same points, and zooming in only ever ADDS points: +a point drawn while zoomed out is still drawn after you zoom into it. A fresh +``random.sample`` per request would make the scene boil under the smallest +camera move, which reads as noise and destroys any sense that you are looking at +a stable structure. + +**A request costs about what it returns, not the size of the dataset.** That +order, and a spatial grid over it, are built once per data change into a +:class:`ProjectionIndex` -- on a worker thread, while the previous one keeps +answering (see :class:`ProjectionCache`). A view query then reads a few cells or +a short prefix of the order instead of every sample. Recomputing the sample +over the whole frame on every request took 18-45 s at 3M samples, which is +what this layout exists to avoid. +""" + +from __future__ import annotations + +import hashlib +import logging +import os +import threading +import time +import uuid + +import numpy as np +import pandas as pd + +from weightslab.proto import experiment_service_pb2 as pb2 + +logger = logging.getLogger(__name__) + +DEFAULT_PREFIX = "umap" +# How many points one view returns when the client does not ask for a number +# (the board does not, unless a browser overrides it): see default_max_points. +ENV_MAX_POINTS = "WEIGHTSLAB_PROJECTION_MAX_POINTS" +DEFAULT_MAX_POINTS = 70_000 +# Above this the browser's main thread spends longer packing the typed arrays +# than the GPU spends drawing them; the client can ask for less, never more. +HARD_MAX_POINTS = 400_000 +_WARNED_MAX_POINTS: set = set() +AXES = ("x", "y", "z") +# The most distinct values a label column may have and still count as "the +# clusters" for stratification. +MAX_GROUPS = 256 +# Groups no bigger than this are ordered exactly within themselves (see +# priority_order). Bigger ones use the hash rank directly: as good for them, and +# it spares sorting millions of points twice. +SMALL_GROUP = 4096 + +# Below this many points a query is one vectorised pass over all of them -- +# exact, and faster than consulting the cell grid. +CELL_INDEX_MIN_POINTS = 200_000 +# Cells per axis of the spatial grid. Both keep a cell id in 16 bits, which is +# what lets numpy sort by cell with its linear-time radix sort. +GRID_CELLS = {2: 256, 3: 32} +# A query whose candidate cells hold at most this share of the cloud reads those +# cells; a bigger one walks the priority order and stops at the budget, which +# for a region that big is the shorter walk. +CELL_PATH_SHARE = 1 / 8 + + +def default_max_points() -> int: + """How many points a view returns when the client leaves it to the server. + + ``WEIGHTSLAB_PROJECTION_MAX_POINTS`` if set to a positive integer, else + ``DEFAULT_MAX_POINTS`` (70,000); capped at ``HARD_MAX_POINTS``. Read on every + request rather than once, so setting it from the Studio's notebook + (``os.environ[...] = "150000"``) applies at the next camera move. + """ + raw = os.environ.get(ENV_MAX_POINTS, "").strip() + if not raw: + return DEFAULT_MAX_POINTS + try: + value = int(raw) # "150000" and "150_000" both parse + except ValueError: + value = 0 + if value <= 0: + if raw not in _WARNED_MAX_POINTS: # once per value, not per camera move + _WARNED_MAX_POINTS.add(raw) + logger.warning(f"[projection] {ENV_MAX_POINTS}={raw!r} is not a positive " + f"integer; drawing {DEFAULT_MAX_POINTS:,} points per view") + return DEFAULT_MAX_POINTS + return min(value, HARD_MAX_POINTS) + + +def coordinate_columns(frame: pd.DataFrame, prefix: str = DEFAULT_PREFIX) -> list: + """The projection columns present in *frame*, in x, y, z order. + + Accepts both the ``signals//`` spelling the training write-back produces and + a bare ``umap_x``, so a projection written by hand (or restored from an H5 + that flattened the prefix) still renders. + """ + found = [] + for axis in AXES: + for candidate in (f"signals//{prefix}_{axis}", f"{prefix}_{axis}"): + if candidate in frame.columns: + found.append(candidate) + break + else: + break # axes must be contiguous: x, then y, then z + return found + + +def available_prefixes(frame: pd.DataFrame) -> list: + """Every projection stored on *frame*, by prefix. + + A prefix counts as present only when its ``_x`` AND ``_y`` columns both + exist -- a half-written projection would otherwise show up in the picker and + fail the moment it was selected. ``"umap"`` (the live one) is listed first + when present; the rest are alphabetical. + + And only when something registered it as a projection (see + :mod:`weightslab.projection.registry`): a column pair alone is not enough, + or ``center_x`` / ``center_y`` signals turn up as a projection called + ``center``. A run with no registry at all predates it, and keeps the old + column-pair discovery. + """ + found = set() + for column in frame.columns: + name = str(column) + base = name[len("signals//"):] if name.startswith("signals//") else name + if not base.endswith("_x"): + continue + prefix = base[:-2] + if not prefix: + continue + # >= 2, not merely truthy: a prefix with only an _x column would be + # offered in the picker and then render nothing when selected. + if len(coordinate_columns(frame, prefix)) >= 2: + found.add(prefix) + try: + from weightslab.projection.registry import known_prefixes + registered = known_prefixes() + except Exception: + registered = None + if registered is not None: + found = {p for p in found if p in registered or p == DEFAULT_PREFIX} + ordered = sorted(found - {DEFAULT_PREFIX}) + return ([DEFAULT_PREFIX] if DEFAULT_PREFIX in found else []) + ordered + + +def grouping_key(frame: pd.DataFrame, coords: np.ndarray, color_column: str, + extent: tuple | None = None) -> np.ndarray: + """The "cluster" each point belongs to, as integer codes, for stratification. + + Preference order, most to least meaningful: + + 1. the categorical column the user is colouring by -- if they are looking + at classes, classes are the clusters; + 2. ``target`` when it is low-cardinality (the usual supervised label); + 3. spatial cells of the cloud itself -- a label-free fallback that still + guarantees no *visible* region of the projection can vanish, which is + the property that actually matters on screen. + + *frame*'s rows must line up with *coords*. + """ + return _group_codes(frame, None, coords, color_column, extent)[1] + + +def _group_codes(frame, rows, coords, color_column, extent=None): + """``(column, codes)`` for :func:`grouping_key`, over *rows* of *frame* only. + + *column* is the label column used, or ``""`` for the spatial fallback. + """ + for candidate in (color_column, "target", "signals//target"): + name = (candidate or "").strip() + if not name: + continue + series = _as_series(frame, name) + if series is None: + continue + values = series.to_numpy() + if rows is not None: + values = values[rows] + codes = _label_codes(values) + if codes is not None: + return name, codes + return "", _spatial_codes(coords, extent) + + +def _label_codes(values) -> np.ndarray | None: + """Integer codes for a label column, or None unless 1 < distinct <= MAX_GROUPS. + + Probed on a sample first: factorising ten million free-text values only to + learn there are millions of them is the expensive way to say no. NaN is a + group of its own, as the string conversion this replaces made it. + """ + values = np.asarray(values) + n = len(values) + if n == 0: + return None + try: + if n > 100_000: + probe = values[:: n // 100_000] + if len(pd.unique(probe)) > MAX_GROUPS: + return None + codes, uniques = pd.factorize(values, use_na_sentinel=False) + except TypeError: + # Unhashable cells (lists, dicts): compare them by their text. + try: + codes, uniques = pd.factorize(values.astype(str), use_na_sentinel=False) + except Exception: + return None + if not 1 < len(uniques) <= MAX_GROUPS: + return None + return codes.astype(np.int32, copy=False) + + +def _spatial_codes(coords, extent=None) -> np.ndarray: + """Cell of an 8^d grid for each point: the label-free grouping. + + Anchored to the WHOLE cloud's extent, never to the current box. A grid + recomputed per request re-labels every point as the camera moves, so the + chosen subset changes wholesale and "zooming in only adds points" stops + holding -- the cloud reshuffles under the user instead of gaining detail. + """ + cells = 8 + coords = np.asarray(coords, dtype=np.float64) + if coords.shape[0] == 0: + return np.zeros(0, dtype=np.int32) + if extent is not None: + lo, hi = np.asarray(extent[0], dtype=float), np.asarray(extent[1], dtype=float) + else: + lo, hi = coords.min(axis=0), coords.max(axis=0) + span = np.where((hi - lo) > 0, hi - lo, 1.0) + idx = np.clip(((coords - lo) / span * cells).astype(np.int64), 0, cells - 1) + weights = cells ** np.arange(idx.shape[1] - 1, -1, -1) + return (idx * weights).sum(axis=1).astype(np.int32) + + +def priority_order(rank: np.ndarray, groups: np.ndarray) -> np.ndarray: + """Every point's index, in the order points are drawn. + + A request draws the first N points of this order that are inside its + region, which is what makes the sample stable: a fixed order means the same + request returns the same points, a bigger budget only adds points, and so + does a smaller region (zooming in). + + The order is stratified by group, by the square root of the group's size. + A flat sample of 50k points out of 10M gives a 2k-point cluster ~10 dots -- + it reads as noise, or vanishes, and the user concludes the model has no such + cluster. Ordering each point by ``position_in_its_group / sqrt(group_size)`` + draws a group of n points at a density proportional to ``1/sqrt(n)``: a + cluster 100x smaller is drawn 10x denser than a flat sample would, and every + group, however small, has its first point near the very front. + + This replaces a fixed 10% floor per group. That floor was paid smallest + group first until the budget ran out, so as soon as 10% of the data + exceeded the budget -- 500k samples at the default 50k -- it spent the whole + budget on the two smallest classes and drew the others not at all. + + *rank* is each point's stable hash in [0, 1) (see :func:`_stable_rank`). + Groups up to :data:`SMALL_GROUP` points are ordered exactly by it; for + bigger ones the rank times the group size IS the position, to within noise + that does not matter at that size. + """ + rank = np.asarray(rank, dtype=np.float64) + codes = np.asarray(groups) + if codes.dtype.kind not in "iu": + codes, _ = pd.factorize(codes, use_na_sentinel=False) + sizes = np.bincount(codes) + root = np.sqrt(sizes, dtype=np.float64) + key = rank * root[codes] + small = sizes <= SMALL_GROUP + if small.any(): + members = np.flatnonzero(small[codes]) + order = members[np.lexsort((rank[members], codes[members]))] + group = codes[order] + position = np.arange(len(order)) - np.searchsorted(group, group, side="left") + # The rank term only breaks ties between groups' equal positions (every + # group's first point sits at 0); it is far below the gap between two + # positions of one group. + key[order] = position / root[group] + rank[order] * 1e-6 + # Sorting the keys was the most expensive step of a build: a comparison + # sort of every point, ~2 s at 10M. Quantised to 16 bits they radix-sort in + # linear time instead, points sharing a bucket keeping their row order. Any + # FIXED order keeps the properties above; this one only blurs the + # stratification by 1/65536 of the key range. + top = float(key.max()) if len(key) else 0.0 + if top <= 0: + return np.arange(len(key)) + buckets = np.minimum(key * (65535.0 / top), 65535.0).astype(np.uint16) + return np.argsort(buckets, kind="stable") + + +def stratified_keep(ids: np.ndarray, groups: np.ndarray, budget: int) -> np.ndarray: + """Indices to draw out of *ids*: the first *budget* of :func:`priority_order`, + in input order (so ids and colours stay aligned).""" + n = len(ids) + if n <= budget: + return np.arange(n) + keep = priority_order(_stable_rank(ids), groups)[:budget] + keep.sort() + return keep + + +def _mix_to_unit(x: np.ndarray) -> np.ndarray: + """splitmix64-style finaliser, mapped to [0, 1): deterministic, + process-independent, and well spread for sequential inputs.""" + x = x.astype(np.uint64, copy=True) + x ^= x >> np.uint64(30) + x *= np.uint64(0xBF58476D1CE4E5B9) + x ^= x >> np.uint64(27) + x *= np.uint64(0x94D049BB133111EB) + x ^= x >> np.uint64(31) + return (x >> np.uint64(11)).astype(np.float64) / float(1 << 53) + + +def _stable_rank(sample_ids: np.ndarray) -> np.ndarray: + """Deterministic [0, 1) rank per sample id -- the stable-subsample key. + + Never Python's ``hash``: it is salted per process, so the cloud would + resample itself every time the server restarted. + """ + ids = np.asarray(sample_ids) + + # Fast path: WeightsLab sample ids are integers (often carried as strings). + # A vectorised integer mix is ~100x quicker than hashing text. + as_int = None + if ids.dtype.kind in "iu": + as_int = ids.astype(np.int64, copy=False) + elif ids.dtype.kind == "f": + if np.isfinite(ids).all(): + as_int = ids.astype(np.int64) + else: + try: + as_int = ids.astype(np.int64) + except (TypeError, ValueError, OverflowError): + as_int = None + if as_int is not None: + return _mix_to_unit(as_int) + + # Non-numeric ids (uuids, paths): FNV-1a over the code points, vectorised + # across ids. An md5 per id was correct but a Python loop -- ~20 s at 10M. + text = ids.astype(str) + width = text.dtype.itemsize // 4 + points = text.view(np.uint32).reshape(len(text), width) + h = np.full(len(text), 0xCBF29CE484222325, dtype=np.uint64) + prime = np.uint64(0x100000001B3) + for j in range(width): + code = points[:, j].astype(np.uint64) + # Skip the zero padding: a hash that depended on the array's width + # would change the moment a longer id joined the dataset. + h = np.where(code != 0, (h ^ code) * prime, h) + return _mix_to_unit(h) + + +def _as_series(frame: pd.DataFrame, name: str): + """Series for *name* whether it is a column or an index level (same helper + GetHistogram uses, for the same reason: ``origin`` may live in either).""" + 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 + + +def _index_level(frame: pd.DataFrame, name: str): + """Index level *name*, or ``None``. + + Levels must be addressed BY NAME, never by position. The ledger frame is + indexed ``(sample_id, annotation_id)`` but the data service's view is + indexed ``(origin, sample_id)`` -- so "level 0" is the sample id in one and + the split name in the other. Assuming a position silently hands back + ``'train_loader'`` as a sample id and filters every row away. + """ + names = list(getattr(frame.index, "names", []) or []) + if name in names: + return frame.index.get_level_values(name) + return None + + +def _sample_ids(frame: pd.DataFrame) -> np.ndarray: + """Sample ids for *frame*, whatever its index layout.""" + level = _index_level(frame, "sample_id") + if level is not None: + return level.to_numpy() + index = frame.index + if isinstance(index, pd.MultiIndex): + return index.get_level_values(0).to_numpy() + return index.to_numpy() + + +def _sample_ids_at(frame: pd.DataFrame, rows: np.ndarray) -> np.ndarray: + """Sample ids of *rows* only -- :func:`_sample_ids` without materialising + the level for every row of a 10M-row frame to read a few thousand.""" + index = frame.index + if isinstance(index, pd.MultiIndex): + names = list(index.names) + level = names.index("sample_id") if "sample_id" in names else 0 + codes = np.asarray(index.codes[level])[rows] + values = index.levels[level].take(np.where(codes < 0, 0, codes)).to_numpy() + if (codes < 0).any(): + values = values.astype(object) + values[codes < 0] = np.nan + return values + return index.take(rows).to_numpy() + + +def restrict_to_sample_ids(frame: pd.DataFrame, sample_ids) -> pd.DataFrame: + """Rows of *frame* whose sample id is in *sample_ids* (compared as strings). + + An empty result is returned as such (the caller reports it): silently + serving the whole frame when the ids do not match made a broken restriction + look like a restriction that was never applied. + + Converts every id in *frame* on each call; :class:`SampleIdLookup` is the + cached form the data service uses. + """ + wanted = {str(s) for s in sample_ids} + if not wanted: + return frame + keep = pd.Series(_sample_ids(frame)).astype(str).isin(wanted).to_numpy() + return frame[keep] + + +def _plain_ids(values: pd.Index) -> bool: + """Whether *values* hold plain integer or string ids, which pandas can + hash as they are (anything else is compared as text).""" + return values.dtype.kind in "iu" or ( + values.dtype.kind == "O" and values.inferred_type == "string") + + +def _keys_like(values: pd.Index, sample_ids) -> np.ndarray: + """*sample_ids* as the dtype *values* holds -- the grid sends strings, + a frame may hold integers.""" + ids = np.asarray(sample_ids) + if values.dtype.kind in "iu": + if ids.dtype.kind in "iu": + return ids.astype(np.int64, copy=False) + numbers = pd.to_numeric(pd.Series(ids.astype(str)), errors="coerce").dropna() + return numbers[numbers == numbers.round()].astype(np.int64).to_numpy() + if ids.dtype.kind == "O" and len(ids) and isinstance(ids[0], str): + return ids + return ids.astype(str).astype(object) + + +class SampleIdLookup: + """Rows of a frame by sample id, through the index's own hash tables. + + A grid page is a hundred-odd samples asked for on every poll and scroll; a + lasso selection can be millions. Converting every id of a 10M-row frame to + a string to find them took ~4 s a time. On the data service's + ``(origin, sample_id)`` view the ids are looked up in the sample_id LEVEL + instead -- each distinct id once, through the hash table pandas builds once + and keeps on the level, which every filtered slice of the view shares -- + and the rows are then found by integer level code. A frame indexed by + sample id alone uses its index's table the same way. Only ids that are not + plain integers or strings are converted to text, once per index. + """ + + def __init__(self): + self._lock = threading.Lock() + self._index_obj = None + self._lookup: pd.Index | None = None + + def rows(self, frame: pd.DataFrame, sample_ids) -> np.ndarray: + """Sorted row positions of *frame* holding any of *sample_ids*.""" + index = frame.index + if isinstance(index, pd.MultiIndex): + names = list(index.names) + level = names.index("sample_id") if "sample_id" in names else 0 + values = index.levels[level] + if _plain_ids(values): + wanted = values.get_indexer(_keys_like(values, sample_ids)) # levels are unique + wanted = wanted[wanted >= 0] + if not len(wanted): + return np.zeros(0, dtype=np.int64) + return np.flatnonzero(np.isin(np.asarray(index.codes[level]), wanted)) + elif _plain_ids(index): + keys = _keys_like(index, sample_ids) + found = (index.get_indexer(keys) if index.is_unique + else index.get_indexer_non_unique(keys)[0]) + return _sorted_unique(found[found >= 0]) + with self._lock: + if self._index_obj is not index: + self._lookup = pd.Index(_sample_ids(frame).astype(str)) + self._index_obj = index + lookup = self._lookup + keys = np.asarray(sample_ids).astype(str).astype(object) + found = (lookup.get_indexer(keys) if lookup.is_unique + else lookup.get_indexer_non_unique(keys)[0]) + return _sorted_unique(found[found >= 0]) + + +def _unique_ids_at(frame: pd.DataFrame, rows: np.ndarray) -> np.ndarray: + """The distinct sample ids of *rows* -- de-duplicated by integer level code + rather than by hashing millions of id strings.""" + index = frame.index + if isinstance(index, pd.MultiIndex): + names = list(index.names) + level = names.index("sample_id") if "sample_id" in names else 0 + codes = _sorted_unique(np.asarray(index.codes[level])[rows]) + return index.levels[level].take(codes[codes >= 0]).to_numpy() + return pd.unique(index.take(rows).to_numpy()) + + +def _sorted_unique(values: np.ndarray) -> np.ndarray: + """``np.unique`` for integers, by sort-and-compare: numpy 2's hash-based + unique took 2.2 s on the 1.5M row numbers of one lasso, this ~50 ms.""" + values = np.sort(values) + if len(values) < 2: + return values + return values[np.concatenate(([True], values[1:] != values[:-1]))] + + +def missing_reason(prefix: str, prefixes: list) -> str: + """Why there is nothing to draw for *prefix* -- the text the board's + "No projection found" ribbon shows, so it has to say what to DO: wait for + the first fit, look at the live projection's failure, or plug one in. + """ + head = f"no projection for prefix '{prefix}'" + if prefixes: + return f"{head}. Available: {', '.join(prefixes)}." + # Registered but not on the board yet: the data view is rebuilt in the + # background (at most every few seconds), so a projection written a moment + # ago -- the user's own t-SNE, just plugged in -- is not in it yet. Saying + # "plug one in" right after they did would send them looking for a bug. + try: + from weightslab.projection.registry import known_prefixes + just_written = sorted(known_prefixes() or ()) + except Exception: + just_written = [] + if just_written: + were = "was" if len(just_written) == 1 else "were" + return (f"{head} on the board yet: {', '.join(just_written)} {were} just " + f"written and should appear within a few seconds.") + try: + from weightslab.projection import get_tracker + tracker = get_tracker() + except Exception: + tracker = None + if tracker is not None and tracker.signal_prefix == prefix: + if tracker.disabled_reason: + return f"{head}: the live projection stopped after {tracker.disabled_reason}." + mismatch = getattr(tracker, "last_mismatch", None) + if mismatch and not tracker.steps_trained: + rows, samples = mismatch + return (f"{head}: the hooked layer {tracker.layer_name!r} gives {rows} " + f"feature rows for a batch of {samples} samples, not one per " + f"sample. Point it at a per-sample layer with " + f"projection={{'layer': ...}} in wl.watch_or_edit(model, ...).") + if tracker._handle is not None: + return (f"{head} yet: the live UMAP places its first points after its " + f"first fit (every {tracker.every_n_steps} training steps).") + return (f"{head}: the built-in projection is off for this run. Plug in your " + f"own with wl.save_projection_coords(coords, batch_ids, prefix=...).") + + +# --- the view index ------------------------------------------------------------ + + +class ViewRegion: + """The part of the cloud one request asks for: a box, optionally cut by the + camera's frustum planes (a point is inside when ``n.p + d >= 0`` for each). + """ + + def __init__(self, lo, hi, normals=None, offsets=None): + # float32, like the coordinates: a float64 bound would widen every + # coordinate it is compared with. + self.lo = np.asarray(lo, dtype=np.float32) + self.hi = np.asarray(hi, dtype=np.float32) + self.normals = None if normals is None else np.asarray(normals, dtype=np.float32) + self.offsets = None if offsets is None else np.asarray(offsets, dtype=np.float32) + + @classmethod + def from_request(cls, request, dims: int) -> "ViewRegion | None": + """The request's region, or None for the whole cloud.""" + planes = np.asarray(list(getattr(request, "frustum_planes", []) or []), + dtype=np.float64) + normals = offsets = None + if planes.size and planes.size % 4 == 0: + planes = planes.reshape(-1, 4) + planes = planes[np.isfinite(planes).all(axis=1)] + if len(planes): + # A 2-D projection lies on z = 0, so the plane's z term drops. + normals, offsets = planes[:, :dims], planes[:, 3] + if not request.has_bounds and normals is None: + return None + if request.has_bounds: + lo = np.array([request.min_x, request.min_y, request.min_z][:dims], dtype=np.float64) + hi = np.array([request.max_x, request.max_y, request.max_z][:dims], dtype=np.float64) + # A NaN bound is no bound (and must mean the same on every path). + lo = np.where(np.isnan(lo), -np.inf, lo) + hi = np.where(np.isnan(hi), np.inf, hi) + else: + lo = np.full(dims, -np.inf) + hi = np.full(dims, np.inf) + return cls(lo, hi, normals, offsets) + + def contains(self, xyz: np.ndarray) -> np.ndarray: + """Boolean mask of the rows of *xyz* inside the region.""" + inside = np.ones(len(xyz), dtype=bool) + for axis in range(xyz.shape[1]): + if np.isfinite(self.lo[axis]): + inside &= xyz[:, axis] >= self.lo[axis] + if np.isfinite(self.hi[axis]): + inside &= xyz[:, axis] <= self.hi[axis] + if self.normals is not None and inside.any(): + # Planes only on what the box kept: one small matmul instead of six + # passes over every candidate. + idx = np.flatnonzero(inside) + ok = (xyz[idx] @ self.normals.T + self.offsets >= 0).all(axis=1) + inside[idx[~ok]] = False + return inside + + +# The lasso is rasterised at canvas resolution, capped to this many pixels on +# the long side (a 4K canvas is the most anyone draws a loop on). +MAX_LASSO_RASTER = 4096 + + +def points_in_lasso(xyz: np.ndarray, view_projection, polygon: np.ndarray, + viewport) -> np.ndarray: + """Which rows of *xyz* land inside *polygon* on screen. + + The same projection as the viewer's own ``idsInPolygon``: *view_projection* + is three.js's ``projectionMatrix * matrixWorldInverse`` (column-major, as + its ``elements``), *polygon* is in canvas pixels with y down, and *viewport* + is the canvas size. Points behind the camera are never inside. + + The loop is rasterised once with PIL and each point is then a single mask + lookup -- an even-odd test against every edge would cost points x edges, + which for a few million points and a hand-drawn loop is seconds. + """ + from PIL import Image, ImageDraw + + result = np.zeros(len(xyz), dtype=bool) + width, height = float(viewport[0]), float(viewport[1]) + if width <= 0 or height <= 0 or len(polygon) < 3 or not len(xyz): + return result + matrix = np.asarray(view_projection, dtype=np.float64).reshape(4, 4).T + dims = xyz.shape[1] + # A 2-D projection lies on z = 0, so the matrix's z column drops out. + clip = np.asarray(xyz, dtype=np.float64) @ matrix[:, :dims].T + matrix[:, 3] + front = np.flatnonzero(clip[:, 3] > 1e-9) + if not len(front): + return result + w = clip[front, 3] + scale = min(1.0, MAX_LASSO_RASTER / max(width, height)) + cols, rows = max(1, int(round(width * scale))), max(1, int(round(height * scale))) + px = (clip[front, 0] / w + 1.0) * 0.5 * cols + py = (1.0 - clip[front, 1] / w) * 0.5 * rows + on_canvas = np.flatnonzero((px >= 0) & (px < cols) & (py >= 0) & (py < rows)) + + image = Image.new("1", (cols, rows), 0) + ImageDraw.Draw(image).polygon( + [(float(x) * scale, float(y) * scale) for x, y in polygon], fill=1, outline=1) + mask = np.asarray(image, dtype=bool) + hit = mask[py[on_canvas].astype(np.int64), px[on_canvas].astype(np.int64)] + result[front[on_canvas[hit]]] = True + return result + + +class _Cells: + """A uniform grid over the cloud: each cell's points, in priority order.""" + + __slots__ = ("g", "lo", "size", "start", "prio", "xyz") + + +def _build_cells(xyz: np.ndarray, lo: np.ndarray, hi: np.ndarray) -> _Cells: + """Group *xyz* (already in priority order) by grid cell. + + A stable sort by cell keeps each cell's points in priority order, so a cell + is a sorted run of priorities -- the property the cell path of + :meth:`ProjectionIndex.select` relies on. + """ + dims = xyz.shape[1] + g = GRID_CELLS[dims] + lo = np.asarray(lo, dtype=np.float64) + span = np.where((np.asarray(hi) - lo) > 0, np.asarray(hi) - lo, 1.0) + size = span / g + cell = np.zeros(len(xyz), dtype=np.uint16) # 32^3 and 256^2 both fit: radix sort + for axis in range(dims): + part = ((xyz[:, axis] - np.float32(lo[axis])) * np.float32(1.0 / size[axis])).astype(np.int32) + np.clip(part, 0, g - 1, out=part) + cell *= g + cell += part.astype(np.uint16) + order = np.argsort(cell, kind="stable") + start = np.zeros(g ** dims + 1, dtype=np.int64) + np.cumsum(np.bincount(cell, minlength=g ** dims), out=start[1:]) + cells = _Cells() + cells.g, cells.lo, cells.size, cells.start = g, lo, size, start + cells.prio = order.astype(np.int32 if len(order) < 2 ** 31 else np.int64) + cells.xyz = np.take(xyz, order, axis=0) + return cells + + +def _column_values(frame: pd.DataFrame, name: str) -> np.ndarray: + """*name* as a float array; NaN where it is not a number. No copy when the + column already is float (the signals columns are, see DataService).""" + column = frame[name] + values = column.to_numpy() + if values.dtype.kind == "f": + return values + return pd.to_numeric(column, errors="coerce").to_numpy(dtype=np.float64) + + +def _column_extent(columns): + """``(lo, hi)`` over 1-D coordinate columns. Reduced per column: the same + min/max along axis 0 of the stacked (n, 3) array took ~1 s at 10M rows, + contiguous columns ~60 ms.""" + lo = np.array([float(c.min()) for c in columns], dtype=np.float32) + hi = np.array([float(c.max()) for c in columns], dtype=np.float32) + return lo, hi + + +def _sample_level_mask(frame: pd.DataFrame): + """Sample-level rows only, or None for "all of them". + + On a detection/segmentation frame the instance rows (annotation_id >= 1) + carry no coordinates and would come back as a wall of NaN. + + Addressed by NAME: the data service's view has no annotation level at all + (it is indexed by origin/sample_id), so a positional "level 1 == 0" test + there compares sample ids against 0 and throws every row away. + + Values also arrive as int64 on some frames and object/strings on others, + so coerce before comparing; and never let this filter empty the frame, + which would report "no finite coordinates yet" on a working projection. + """ + annotations = _index_level(frame, "annotation_id") + if annotations is None: + return None + try: + level = pd.to_numeric(pd.Series(annotations), errors="coerce").to_numpy() + keep = level == 0 + return keep if keep.any() else None + except Exception: + return None + + +def _frame_rows(frame: pd.DataFrame, rows: np.ndarray, names) -> pd.DataFrame: + """*rows* of *frame*, with only the columns in *names* -- the colour, split + and discarded fills read a few columns of the drawn points, and slicing a + wide frame for them cost more than the rest of the response. Built column by + column (each keeps its dtype) rather than through ``iloc``.""" + data = {} + for name in dict.fromkeys(n for n in names if n): + if name in frame.columns: + column = frame[name] + if isinstance(column, pd.Series): + data[name] = column.array.take(rows) + return pd.DataFrame(data, index=frame.index.take(rows)) + + +def _fill_origins_from_index(response, frame: pd.DataFrame, rows: np.ndarray) -> bool: + """:func:`_fill_origins` straight from the index's level codes, when the + split is an index level (the data service's view) -- no string per point. + False when it is not, for the general path to handle.""" + index = frame.index + if "origin" in frame.columns or not isinstance(index, pd.MultiIndex) \ + or "origin" not in index.names: + return False + level = list(index.names).index("origin") + codes = np.asarray(index.codes[level])[rows] + used, inverse = np.unique(codes, return_inverse=True) + names = index.levels[level] + response.origins.extend([str(names[c]) if c >= 0 else "nan" for c in used]) + response.origin_ids.extend(np.asarray(inverse, dtype=np.int32).tolist()) + return True + + +class ProjectionIndex: + """One projection of one data frame, laid out so a view query costs about + what it returns. + + Holds every projected point in :func:`priority_order` (``rows`` are their + rows in *frame*, ``xyz`` their coordinates when the index was built), and, + from :data:`CELL_INDEX_MIN_POINTS` points up, the same points grouped by a + uniform grid. :meth:`select` answers "the first N points of the order inside + this region" from whichever of the two reads less. + + What is drawn is read from *frame* at answer time -- coordinates, colours, + split, discarded -- so a point's position and colour are live even while + the index that chose it is a few seconds old. Only WHICH points are chosen + waits for the next build. + """ + + def __init__(self, *, frame, prefix, columns, rows, xyz, lo, hi, cells, + color_column, group_column, fingerprint, build_seconds): + self.frame = frame + self.prefix = prefix + self.columns = columns + self.dims = len(columns) + self.rows = rows + self.xyz = xyz + self.lo = lo + self.hi = hi + self.cells = cells + self.color_column = color_column + self.group_column = group_column + self.fingerprint = fingerprint + self.build_seconds = build_seconds + + @property + def n(self) -> int: + return len(self.rows) + + def ensure_cells(self) -> None: + """Build the cell grid if it was deferred (see :meth:`build`).""" + if self.cells is None and self.n >= CELL_INDEX_MIN_POINTS: + self.cells = _build_cells(self.xyz, self.lo, self.hi) + + @classmethod + def build(cls, frame, prefix: str = "", color_column: str = "", previous=None, + defer_cells: bool = False): + """``(index, None)``, or ``(None, failure_response)`` when there is + nothing to index. + + With *previous* (an index of the same frame), unchanged data returns + *previous* itself, and moved coordinates reuse its priority order -- + the hash, grouping and sort are what a full build spends its time on, + and none of them depend on where the points are. + + *defer_cells* leaves the cell grid to :meth:`ensure_cells`: queries + are exact without it, only slower, so the first answer need not wait. + """ + started = time.perf_counter() + if frame is None or len(frame) == 0: + return None, pb2.ProjectionResponse( + success=False, message="no data view available") + + # Listed even on the failure paths below, so a user who asked for a prefix + # that isn't there is told which ones ARE -- the common case right after an + # offline re-projection under a new name. + prefixes = available_prefixes(frame) + # No preference from the client means "whatever there is": the live + # projection when it exists, else the first stored one. Defaulting blindly + # to "umap" failed the first request of every run that turned the built-in + # off and brought its own t-SNE, and the board only recovered a poll later. + asked = (prefix or "").strip() + prefix = asked or (DEFAULT_PREFIX if DEFAULT_PREFIX in prefixes or not prefixes + else prefixes[0]) + + columns = coordinate_columns(frame, prefix) + if len(columns) < 2: + return None, pb2.ProjectionResponse( + success=False, + message=missing_reason(prefix, prefixes), + available_prefixes=prefixes, + ) + + raw = [_column_values(frame, c) for c in columns] + keep = np.ones(len(frame), dtype=bool) + for values in raw: + keep &= np.isfinite(values) + level = _sample_level_mask(frame) + if level is not None: + keep &= level + if not keep.any(): + return None, pb2.ProjectionResponse( + success=False, + message=f"projection '{prefix}' has no coordinates yet", + total_available=0, + available_prefixes=prefixes, + ) + + color_column = (color_column or "").strip() + fingerprint = ( + len(frame), + hashlib.blake2b(np.packbits(keep).tobytes(), digest_size=16).digest(), + tuple(float(np.nansum(v, dtype=np.float64)) for v in raw), + ) + rows = group_column = None + if (previous is not None and previous.frame is frame and previous.prefix == prefix + and previous.columns == columns and previous.color_column == color_column): + if previous.fingerprint == fingerprint: + return previous, None + if previous.fingerprint[:2] == fingerprint[:2]: + # Same samples, new coordinates. + rows, group_column = previous.rows, previous.group_column + + row_type = np.int32 if len(frame) < 2 ** 31 else np.int64 + if rows is None: + # Every row projected is the common case: skip the gathers. + everything = bool(keep.all()) + rows_asc = np.arange(len(frame), dtype=row_type) if everything else np.flatnonzero(keep) + columns_xyz = [values if everything else values[rows_asc] for values in raw] + lo, hi = _column_extent(columns_xyz) + xyz = np.stack(columns_xyz, axis=1).astype(np.float32, copy=False) + ids = _sample_ids(frame) + group_column, codes = _group_codes(frame, None if everything else rows_asc, + xyz, color_column, (lo, hi)) + order = priority_order(_stable_rank(ids if everything else ids[rows_asc]), codes) + rows = rows_asc[order].astype(row_type, copy=False) + xyz = np.take(xyz, order, axis=0) + else: + columns_xyz = [values[rows] for values in raw] + lo, hi = _column_extent(columns_xyz) + xyz = np.stack(columns_xyz, axis=1).astype(np.float32, copy=False) + + cells = (_build_cells(xyz, lo, hi) + if len(rows) >= CELL_INDEX_MIN_POINTS and not defer_cells else None) + return cls( + frame=frame, prefix=prefix, columns=columns, rows=rows, xyz=xyz, + lo=lo.astype(np.float64), hi=hi.astype(np.float64), cells=cells, + color_column=color_column, group_column=group_column, + fingerprint=fingerprint, build_seconds=time.perf_counter() - started, + ), None + + # ------------------------------------------------------------------ query + def select(self, region: ViewRegion | None, budget: int): + """``(positions, in_view)``: the first *budget* points of the priority + order inside *region*, as ascending positions in that order, and how + many points the region holds (estimated when the walk stopped early). + + Every path returns exactly the same points for the same region; they + differ only in how much they read. + """ + n = self.n + if region is None: + return np.arange(min(n, budget)), n + if self.cells is None: + hits = np.flatnonzero(region.contains(self.xyz)) + return hits[:budget], len(hits) + starts, ends = self._cell_runs(region) + candidates = int((ends - starts).sum()) + if candidates == 0: + return np.zeros(0, dtype=np.int64), 0 + if candidates <= max(budget, n * CELL_PATH_SHARE): + return self._select_from_cells(region, starts, ends, budget) + return self._select_by_walk(region, budget) + + def _cell_runs(self, region: ViewRegion): + """Start/end offsets of the grid cells the region can reach: those its + box overlaps, minus those wholly outside one of its planes. + + Culling cells against the frustum is what keeps an oblique view cheap: + its bounding box takes in most of the cloud, the frustum itself a + slice of it.""" + cells = self.cells + g = cells.g + axes = [] + for axis in range(self.dims): + lo, hi = float(region.lo[axis]), float(region.hi[axis]) + if lo == np.inf or hi == -np.inf: + empty = np.zeros(0, dtype=np.int64) + return empty, empty + first = 0 if lo == -np.inf else int(np.floor((lo - cells.lo[axis]) / cells.size[axis])) + last = g - 1 if hi == np.inf else int(np.floor((hi - cells.lo[axis]) / cells.size[axis])) + first, last = max(first, 0), min(last, g - 1) + if first > last: + empty = np.zeros(0, dtype=np.int64) + return empty, empty + axes.append(np.arange(first, last + 1)) + index = [grid.ravel() for grid in np.meshgrid(*axes, indexing="ij")] + ids = index[0] + for part in index[1:]: + ids = ids * g + part + if region.normals is not None: + keep = np.ones(len(ids), dtype=bool) + corner = [cells.lo[a] + index[a] * cells.size[a] for a in range(self.dims)] + for normal, offset in zip(region.normals.astype(np.float64), region.offsets): + # The cell's corner furthest along the normal: if even that one + # is outside the plane, the whole cell is. + reach = float(offset) + sum( + normal[a] * (corner[a] + (cells.size[a] if normal[a] > 0 else 0.0)) + for a in range(self.dims)) + keep &= reach >= 0 + ids = ids[keep] + return cells.start[ids], cells.start[ids + 1] + + @staticmethod + def _cell_positions(starts, ends) -> np.ndarray: + """Every offset of the runs ``[start, end)``, concatenated.""" + lengths = ends - starts + live = lengths > 0 + starts, lengths = starts[live], lengths[live] + total = int(lengths.sum()) + before = np.concatenate(([0], np.cumsum(lengths)[:-1])) + return np.repeat(starts - before, lengths) + np.arange(total) + + def select_all(self, region: ViewRegion | None) -> np.ndarray: + """Ascending positions of EVERY point inside *region* -- no budget. + What a selection needs: the drawn points are only a subsample.""" + n = self.n + if region is None: + return np.arange(n) + if self.cells is not None: + starts, ends = self._cell_runs(region) + candidates = int((ends - starts).sum()) + if candidates == 0: + return np.zeros(0, dtype=np.int64) + if candidates <= n // 2: + positions = self._cell_positions(starts, ends) + inside = region.contains(self.cells.xyz[positions]) + found = self.cells.prio[positions[inside]].astype(np.int64) + found.sort() + return found + return np.flatnonzero(region.contains(self.xyz)) + + def lasso(self, request) -> np.ndarray: + """Ascending positions of every point whose screen position, under the + request's camera, falls inside its lasso polygon -- whether or not the + view ever drew it. The request's box and planes (the loop's own + sub-frustum) narrow the candidates first.""" + polygon = np.asarray(list(request.lasso_px), dtype=np.float64) + viewport = list(request.lasso_viewport) + matrix = list(request.view_projection) + if polygon.size < 6 or polygon.size % 2 or len(viewport) != 2 or len(matrix) != 16: + return np.zeros(0, dtype=np.int64) + candidates = self.select_all(ViewRegion.from_request(request, self.dims)) + if not len(candidates): + return candidates + inside = points_in_lasso(self.xyz[candidates], matrix, polygon.reshape(-1, 2), viewport) + return candidates[inside] + + def _select_from_cells(self, region, starts, ends, budget): + cells = self.cells + positions = self._cell_positions(starts, ends) + inside = region.contains(cells.xyz[positions]) + prios = cells.prio[positions[inside]].astype(np.int64) + in_view = len(prios) + if in_view > budget: + prios = np.partition(prios, budget - 1)[:budget] + prios.sort() + return prios, in_view + + def _select_by_walk(self, region, budget): + """Walk the priority order in growing chunks until *budget* points of + the region have been seen. For a region holding a share s of the cloud + that reads about budget / s points.""" + n = self.n + found, hits = 0, [] + start, step = 0, max(4 * budget, 65_536) + while start < n and found < budget: + stop = min(n, start + step) + idx = np.flatnonzero(region.contains(self.xyz[start:stop])) + start + hits.append(idx) + found += len(idx) + start, step = stop, step * 2 + positions = np.concatenate(hits)[:budget] if hits else np.zeros(0, dtype=np.int64) + in_view = found if start >= n else max(found, int(round(found * n / start))) + return positions, in_view + + # --------------------------------------------------------------- response + def _coords_at(self, rows: np.ndarray, positions: np.ndarray) -> np.ndarray: + """Live coordinates of *rows*, falling back to the indexed ones for any + row that no longer reads as a number.""" + indexed = self.xyz[positions] + try: + live = np.stack([np.asarray(_column_values(self.frame, c)[rows], dtype=np.float32) + for c in self.columns], axis=1) + except Exception: + return indexed + stale = ~np.isfinite(live).all(axis=1) + if stale.any(): + live[stale] = indexed[stale] + return live + + def respond(self, request, selection: "_Selection | None" = None) -> pb2.ProjectionResponse: + """Answer *request* from this index. *selection* (the current lasso + selection) flags the returned points that belong to it.""" + budget = int(request.max_points) or default_max_points() + budget = max(1, min(budget, HARD_MAX_POINTS)) + positions, in_view = self.select(ViewRegion.from_request(request, self.dims), budget) + rows = self.rows[positions] + frame = self.frame + ids = np.asarray(_sample_ids_at(frame, rows)) + + response = pb2.ProjectionResponse( + success=True, + message=f"{len(rows)} of {in_view} points in view", + sample_ids=ids.astype(str).tolist(), + coords=self._coords_at(rows, positions).reshape(-1).tolist(), + returned=len(rows), + total_in_view=in_view, + total_available=self.n, + dims=self.dims, + available_prefixes=available_prefixes(frame), + ) + # Extent of the WHOLE projection, not of the region: the client frames + # the full cloud from this on first load and fits its grid walls to it. + response.extent_min_x, response.extent_min_y = float(self.lo[0]), float(self.lo[1]) + response.extent_max_x, response.extent_max_y = float(self.hi[0]), float(self.hi[1]) + if self.dims > 2: + response.extent_min_z = float(self.lo[2]) + response.extent_max_z = float(self.hi[2]) + + sub = _frame_rows(frame, rows, ("origin", "discarded", request.color_column, + "target", "signals//target")) + if not _fill_origins_from_index(response, frame, rows): + _fill_origins(response, sub) + _fill_colors(response, sub, request.color_column) + _fill_discarded(response, sub) + if selection is not None and len(selection.ids) and len(ids): + response.in_selection.extend(selection.flags(frame, rows, ids).tolist()) + + try: + from weightslab.projection import get_tracker + tracker = get_tracker() + # Only the projection the live encoder writes has fits. A user's own + # t-SNE shown with the encoder's count would claim training it never had. + if tracker is not None and tracker.signal_prefix == self.prefix: + response.fits = int(tracker.steps_trained) + except Exception: + pass + response.layer, response.layer_detail = projection_layer(self.prefix) + return response + + +def projection_layer(prefix: str) -> tuple: + """``(layer, layer_detail)`` that *prefix*'s coordinates were computed from, + ``("", "")`` when unknown (a projection the user computed themselves). + + The live tracker answers for its own prefix -- it may have re-attached to + another layer since it registered; the registry for everything else. + """ + try: + from weightslab.projection import get_tracker + tracker = get_tracker() + if (tracker is not None and tracker.signal_prefix == prefix + and tracker.layer_name): + return tracker.layer_name, tracker.layer_detail or "" + except Exception: + pass + try: + from weightslab.projection.registry import prefix_info + info = prefix_info(prefix) + except Exception: + info = {} + return info.get("layer", ""), info.get("layer_detail", "") + + +def build_projection_response(frame, request) -> pb2.ProjectionResponse: + """Answer one :class:`ProjectionRequest` against *frame*, indexing it on + the spot. + + The data service keeps the index between requests instead (see + :class:`ProjectionCache`); this uncached form is for tests, examples and + one-off callers, and returns exactly what the cached one would. + + Index layout is not assumed. The ledger frame is indexed + ``(sample_id, annotation_id)`` and the data service's view is indexed + ``(origin, sample_id)``, so every level is resolved by NAME -- reading + "level 0" gives a split name on one of them. + """ + index, failure = ProjectionIndex.build(frame, request.prefix, request.color_column) + if failure is not None: + return failure + return index.respond(request) + + +def _respond_in_full(frame, request) -> pb2.ProjectionResponse: + """Every point of *frame*: no region, no budget. For a grid page, which is + small enough to draw whole and is exactly what the user asked to see.""" + plain = pb2.ProjectionRequest() + plain.CopyFrom(request) + plain.has_bounds = False + plain.ClearField("frustum_planes") + plain.max_points = HARD_MAX_POINTS + return build_projection_response(frame, plain) + + +class _Selection: + """A lasso selection: its distinct sample ids, and the rows of the frame it + was drawn on. Rows make the common cases integer work -- applying it to + that same frame, flagging that frame's drawn points, combining it with the + next loop -- and the ids carry it to any other frame.""" + + __slots__ = ("ids", "frame", "rows", "_index") + + def __init__(self, ids: np.ndarray, frame, rows: np.ndarray): + self.ids, self.frame, self.rows = ids, frame, rows + self._index = None + + def rows_in(self, frame, lookup: SampleIdLookup) -> np.ndarray: + """Sorted rows of *frame* holding the selection.""" + if frame is self.frame: + return self.rows + return lookup.rows(frame, self.ids) if len(self.ids) else np.zeros(0, dtype=np.int64) + + def flags(self, frame, rows: np.ndarray, ids: np.ndarray) -> np.ndarray: + """Per point (given by its row of *frame* and its id): selected?""" + if frame is self.frame: + return np.isin(rows, self.rows) + if self._index is None: + self._index = pd.Index(self.ids) + return self._index.get_indexer(ids) >= 0 + + +class _Entry: + __slots__ = ("index", "kind", "checked_at", "pulled_at") + + def __init__(self, index, kind, pulled_at=0.0): + self.index = index + self.kind = kind + self.checked_at = time.monotonic() + self.pulled_at = pulled_at + + +class ProjectionCache: + """The :class:`ProjectionIndex` of each projection being looked at, kept + fresh in the background. + + The first request for a projection builds its index -- the one wait anybody + sees, a few seconds at 10M samples. From then on a request only QUERIES an + index. When the data has moved on (training writes coordinates all the + time), a newer index is built on a worker thread and swapped in while the + old one keeps answering, so no camera move ever waits for a rebuild and the + choice of points is at most one rebuild behind. + + The heavy work also happens outside the data service's lock, which the old + path held for the whole computation -- stalling the grid and the plots for + as long as a projection request took. + """ + + #: Least seconds between two freshness checks of one index -- more when a + #: build takes longer, so a worker spends at most a fifth of its time + #: rebuilding (training keeps the rest). + REFRESH_SECONDS = 5.0 + #: Least seconds between two full pulls of the dataset -- needed only while + #: the data view is filtered and the board still shows everything. + FULL_PULL_SECONDS = 60.0 + + #: How many lasso selections are kept for the data board to apply. + KEPT_SELECTIONS = 4 + + def __init__(self): + self._lock = threading.Lock() + self._entries: dict = {} + self._busy: set = set() + self._ids = SampleIdLookup() + # Selections, by token, oldest first; and the latest, which ADD and + # REFINE combine with. + self._selections: dict = {} + self._selection: _Selection | None = None + self._selection_lookup = SampleIdLookup() + + def note_edit(self) -> None: + """Samples were edited (discarded, tagged): what a pulled copy of the + dataset says about them is now stale. + + While the data view is filtered or sorted, the whole-dataset cloud is + drawn from such a copy, re-pulled at most every FULL_PULL_SECONDS -- + so a sample discarded in the grid stayed un-greyed in the projection + for up to a minute. The next request refreshes instead; the refresh + still runs in the background, so the request itself does not wait. + """ + with self._lock: + for entry in self._entries.values(): + entry.pulled_at = 0.0 + entry.checked_at = 0.0 + + def serve(self, request, *, view, is_filtered, refresh_view=None, pull_full=None): + """Answer *request*. + + *view* returns the data board's current frame as it stands; + *refresh_view* brings it up to date first (and may take a lock); + *pull_full* builds the whole, unfiltered dataset (expensive); + *is_filtered* says whether the view is narrowed by a filter. + """ + follow = bool(request.follow_view) + grid_ids = list(request.restrict_sample_ids) if follow else [] + if grid_ids: + return self._serve_grid(view(), request, grid_ids) + entry, failure = self._entry_for(request, (view, is_filtered, refresh_view, pull_full)) + if failure is not None: + return failure + if getattr(request, "lasso_mode", 0): + return self._select(entry.index, request) + with self._lock: + selection = self._selection + return entry.index.respond(request, selection) + + def _entry_for(self, request, sources): + """``(entry, None)`` for the projection *request* looks at, or + ``(None, failure)``. Builds it the first time; refreshes it in the + background after that.""" + view, is_filtered, _refresh_view, _pull_full = sources + # The whole dataset, unless the board follows a FILTERED view: then the + # cloud is that subset (see DataService.GetProjection for why). + kind = "view" if (request.follow_view and is_filtered()) else "full" + key = (kind, (request.prefix or "").strip()) + color = (request.color_column or "").strip() + + with self._lock: + entry = self._entries.get(key) + # A filtered view the index was not built from is a different set of + # samples: answering from it would show the wrong cloud, so rebuild. + if entry is not None and kind == "view" and entry.index.frame is not view(): + entry = None + if entry is None: + frame, pulled = self._source(kind, sources) + # The first answer is the one the user waits for: the grid only + # speeds queries up, so it is built right after, on a worker. + index, failure = ProjectionIndex.build(frame, key[1], color, defer_cells=True) + if failure is not None: + return None, failure + if index.cells is None and index.n >= CELL_INDEX_MIN_POINTS: + threading.Thread(target=index.ensure_cells, name="WL-ProjectionCells", + daemon=True).start() + entry = _Entry(index, kind, pulled_at=time.monotonic() if pulled else 0.0) + with self._lock: + self._entries[key] = entry + else: + self._maybe_refresh(key, entry, color, sources) + return entry, None + + # -------------------------------------------------------------- selection + def _select(self, index: ProjectionIndex, request) -> pb2.ProjectionResponse: + """A lasso: select on the server, keep it, hand back a token. + + Every sample inside the loop is selected, not just the drawn ones -- at + 10M samples the view draws one in two hundred, and a selection of the + drawn points only would be a sample of a selection. That can be + millions of ids, so they stay here and the data board applies them by + token ("@projection_selection "). + """ + frame = index.frame + rows = np.sort(index.rows[index.lasso(request)]).astype(np.int64) + with self._lock: + previous = self._selection + mode = int(request.lasso_mode) + if previous is not None and len(previous.ids) and mode in ( + pb2.PROJECTION_LASSO_ADD, pb2.PROJECTION_LASSO_REFINE): + # Combined as row numbers -- integers, not millions of id strings. + before = previous.rows_in(frame, self._selection_lookup) + rows = (np.union1d(before, rows) if mode == pb2.PROJECTION_LASSO_ADD + else np.intersect1d(before, rows, assume_unique=True)) + # One sample under two splits is one sample selected. + selection = _Selection(_unique_ids_at(frame, rows), frame, rows) + token = uuid.uuid4().hex[:16] + with self._lock: + self._selection = selection + self._selections[token] = selection + while len(self._selections) > self.KEPT_SELECTIONS: + self._selections.pop(next(iter(self._selections))) + ids = selection.ids + return pb2.ProjectionResponse( + success=True, + message=f"{len(ids)} selected", + selected=len(ids), + selection_token=token, + total_available=index.n, + dims=index.dims, + available_prefixes=available_prefixes(index.frame), + ) + + def selection(self, token: str): + """The sample ids a lasso selected, or None once it has expired.""" + with self._lock: + selection = self._selections.get(token) + return None if selection is None else selection.ids + + def selection_rows(self, token: str, frame) -> np.ndarray | None: + """Rows of *frame* a lasso selected, or None once it has expired. Free + when *frame* is the one the lasso was drawn on -- the usual case.""" + with self._lock: + selection = self._selections.get(token) + return None if selection is None else selection.rows_in(frame, self._selection_lookup) + + def clear_selection(self) -> None: + """Forget the selection ADD and REFINE build on (the board was reset).""" + with self._lock: + self._selection = None + + def full_frame(self): + """The frame the whole dataset was last indexed from, or None.""" + with self._lock: + for (kind, _prefix), entry in self._entries.items(): + if kind == "full": + return entry.index.frame + return None + + def rows_of(self, frame, sample_ids) -> np.ndarray: + """Rows of *frame* holding *sample_ids* (see :class:`SampleIdLookup`).""" + return self._selection_lookup.rows(frame, sample_ids) + + def _serve_grid(self, frame, request, grid_ids): + if frame is None or len(frame) == 0: + return pb2.ProjectionResponse(success=False, message="no data view available") + rows = self._ids.rows(frame, grid_ids) + if not len(rows): + logger.warning( + "GetProjection: none of the grid's %d sample ids (e.g. %s) are in " + "the data view", len(grid_ids), grid_ids[:3]) + return pb2.ProjectionResponse( + success=False, + message=(f"None of the {len(grid_ids)} samples in " + "the grid are in the projection data.")) + return _respond_in_full(frame.iloc[rows], request) + + @staticmethod + def _source(kind, sources): + """``(frame, pulled)``: the frame an index of *kind* is built from.""" + view, is_filtered, refresh_view, pull_full = sources + if kind == "full" and is_filtered() and pull_full is not None: + pulled = pull_full() + if pulled is not None and not getattr(pulled, "empty", True): + return pulled, True + if refresh_view is not None: + return refresh_view(), False + return view(), False + + def _maybe_refresh(self, key, entry, color, sources): + wait = max(self.REFRESH_SECONDS, 4.0 * entry.index.build_seconds) + if (time.monotonic() - entry.checked_at < wait + and entry.index.color_column == color): + return + with self._lock: + if key in self._busy: + return + self._busy.add(key) + threading.Thread( + target=self._refresh, args=(key, entry, color, sources), + name="WL-ProjectionIndex", daemon=True, + ).start() + + def _refresh(self, key, entry, color, sources): + try: + view, is_filtered, refresh_view, pull_full = sources + now = time.monotonic() + if entry.kind == "full" and is_filtered(): + # The whole dataset is not the view right now; re-pulling it + # costs seconds and a full copy, so it is done rarely. + if (now - entry.pulled_at < self.FULL_PULL_SECONDS + and entry.index.color_column == color): + entry.checked_at = now + return + frame, pulled = self._source(entry.kind, sources) + index, failure = ProjectionIndex.build(frame, key[1], color, previous=entry.index) + with self._lock: + if failure is not None: + # Gone (prefix removed, data reset): the next request + # rebuilds in the open and reports why. + self._entries.pop(key, None) + elif index is entry.index: + entry.checked_at = time.monotonic() + else: + self._entries[key] = _Entry( + index, entry.kind, + pulled_at=time.monotonic() if pulled else entry.pulled_at) + except Exception: + logger.exception("[projection] background index refresh failed") + entry.checked_at = time.monotonic() + finally: + with self._lock: + self._busy.discard(key) + + +def _fill_origins(response, frame) -> None: + """Split membership, as an id per point plus the name table. + + Interned rather than repeated per point: at a 50k budget the literal string + "train_loader" 50,000 times is most of the payload. + """ + series = _as_series(frame, "origin") + if series is None: + return + values = series.astype(str).to_numpy() + names, ids = np.unique(values, return_inverse=True) + response.origins.extend(names.astype(str).tolist()) + # .tolist() rather than a comprehension: it converts the whole array in C, + # and these run once per POINT on every refresh of the board. + response.origin_ids.extend(np.asarray(ids, dtype=np.int32).tolist()) + + +def _fill_discarded(response, frame) -> None: + """Per-point discarded flag, so the viewer can grey those samples out.""" + series = _as_series(frame, "discarded") + if series is None: + return + try: + from weightslab.trainer.services.data_service import set_flag_mask + mask = set_flag_mask(series) + except Exception: + # NaN in a nullable flag is NOT "discarded" -- the same trap the + # histogram path documents. + mask = series.fillna(False).astype(bool).to_numpy() + response.discarded.extend(np.asarray(mask, dtype=bool).tolist()) + + +def _text_labels(series: pd.Series): + """``(codes, distinct_texts)`` for *series* -- its values as the text + ``astype(str)`` gives, without converting every value to a string: each + distinct value is converted once.""" + try: + codes, uniques = pd.factorize(series.to_numpy(), use_na_sentinel=False) + texts = pd.Series(uniques).astype(str).to_numpy() + except TypeError: + codes, texts = pd.factorize(series.astype(str).to_numpy(), use_na_sentinel=False) + texts = np.asarray(texts) + # Distinct values may share a text (1 and "1"): merge them as the old + # astype(str)-then-unique did. + merged, remap = np.unique(texts, return_inverse=True) + return remap[codes], merged + + +def _fill_colors(response, frame, column: str) -> None: + """Optional per-point colour source. + + Numeric columns fill ``color_values`` (the client ramps them); anything else + fills ``color_labels`` (the client assigns a categorical palette). Same + numeric-or-not question GetHistogram answers, kept deliberately simple here: + a column that coerces to any finite number is numeric. + """ + column = (column or "").strip() + + # "Colour by split" is the CLIENT's job: it holds the split palette, so a + # sample is the same colour here as in the data grid. Filling labels for it + # would hand the viewer a second, different palette for the same thing. + if column in ("origin", "split"): + return + + if not column: + # Default to the ground-truth label when there is one: on a + # classification set the classes are what the clusters MEAN, and + # colouring by split paints the entire train set one colour, which + # says nothing about the structure on screen. + for candidate in ("target", "signals//target"): + series = _as_series(frame, candidate) + if series is None: + continue + if 1 < len(_text_labels(series)[1]) <= 64: + column = candidate + break + if not column: + return + series = _as_series(frame, column) + if series is None: + return + + # A class id is a label, not a measurement. Running 0..9 through the + # continuous ramp places class 4 "between" 3 and 5 and paints neighbouring + # classes near-identical shades -- the one thing a projection coloured by + # class must not say. A regression target (many distinct values) still + # ramps, which is what it should do. + if column in ("target", "signals//target"): + codes, uniques = _text_labels(series) + if 1 < len(uniques) <= 64: + response.color_labels.extend(uniques[codes].tolist()) + return + + numeric = pd.to_numeric(series, errors="coerce") + if np.isfinite(numeric.to_numpy(dtype=np.float64)).any(): + response.color_values.extend( + numeric.fillna(np.nan).to_numpy(dtype=np.float64).tolist()) + else: + response.color_labels.extend(series.astype(str).tolist()) diff --git a/weightslab/trainer/trainer_services.py b/weightslab/trainer/trainer_services.py index 43e27a6d..9ae007e9 100644 --- a/weightslab/trainer/trainer_services.py +++ b/weightslab/trainer/trainer_services.py @@ -356,6 +356,11 @@ def GetDataSamples(self, request, context): def GetHistogram(self, request, context): logger.debug(f"\nExperimentServiceServicer.GetHistogram({request})") return self._exp_service.data_service.GetHistogram(request, context) + + def GetProjection(self, request, context): + logger.debug(f"\nExperimentServiceServicer.GetProjection({request})") + return self._exp_service.data_service.GetProjection(request, context) + def GetMetaData(self, request, context): logger.debug(f"\nExperimentServiceServicer.GetMetaData({request})") return self._exp_service.data_service.GetMetaData(request, context) diff --git a/weightslab/trainer/trainer_tools.py b/weightslab/trainer/trainer_tools.py index 1b2d467c..b152ae34 100644 --- a/weightslab/trainer/trainer_tools.py +++ b/weightslab/trainer/trainer_tools.py @@ -396,9 +396,9 @@ def load_raw_image(dataset, index): if hasattr(np_img, 'numpy'): np_img = np_img.numpy() if np_img.ndim == 2: - return Image.fromarray(np_img.astype(np.uint8), mode="L") + return Image.fromarray(np_img.astype(np.uint8)) elif np_img.ndim == 3: - return Image.fromarray(np_img.astype(np.uint8), mode="RGB") + return Image.fromarray(np_img.astype(np.uint8)).convert("RGB") else: raise ValueError(f"Unsupported image shape: {np_img.shape}") diff --git a/weightslab/ui/server.py b/weightslab/ui/server.py index a524cb35..3c288139 100644 --- a/weightslab/ui/server.py +++ b/weightslab/ui/server.py @@ -243,6 +243,9 @@ def open_notebook(self, notebooks_dir: str, name: str) -> dict: process = subprocess.Popen( _build_jupyter_cmd(notebooks_dir), cwd=notebooks_dir, + # Never read; inheriting it fails on Windows when the + # parent has none usable (see opencode_process.resolve_or_spawn_opencode). + stdin=subprocess.DEVNULL, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True, @@ -670,6 +673,9 @@ def ensure(self, workspace_dir: str, origin: Optional[str]) -> dict: process = subprocess.Popen( cmd, cwd=workspace_dir, + # Never read; inheriting it fails on Windows when the + # parent has none usable (see opencode_process.resolve_or_spawn_opencode). + stdin=subprocess.DEVNULL, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True, @@ -826,6 +832,9 @@ def _read_example_main(usecase: str) -> Optional[str]: # "/loop 30s ..." hammering the model every few seconds. _LOOP_MIN_INTERVAL_SECONDS = 30.0 +# Best-effort abort/delete when a loop is stopped: short, because the stop +# button waits on it and a wedged OpenCode must not hang the UI. +_LOOP_STOP_TIMEOUT_SECONDS = float(os.environ.get("WL_LOOP_STOP_TIMEOUT_S", "5")) # Caps a single check-in's wall-clock time (_opencode_send_and_collect's own # default is 600s, generous for the SDK/landing chat's own interactive use). # 150s was tried first and was wrong: it was sized for "what's the last loss @@ -1381,6 +1390,17 @@ def _fire(self, job_id: str, base_url: str) -> None: timer.start() def stop(self, job_id: str) -> dict: + """Stop a loop: cancel its next tick AND kill the one it is running now. + + Cancelling the timer only prevents FUTURE ticks. A check-in already + inside _opencode_send_and_collect keeps going until + _LOOP_CHECKIN_TIMEOUT_SECONDS, so closing a loop's tab used to leave + the agent working -- burning tokens on an answer whose tab no longer + exists, and holding the session open. Aborting the session ends that + turn; deleting it releases the session the registry used to leak here. + Both are best-effort with short timeouts: a wedged OpenCode must not + make the stop button hang. + """ with self._lock: job = self._jobs.pop(job_id, None) if job is not None: @@ -1389,6 +1409,18 @@ def stop(self, job_id: str) -> dict: return {"ok": False, "error": f"No loop job {job_id}."} if job.timer is not None: job.timer.cancel() + + if job.session_id: + for path, method in ((f"/session/{job.session_id}/abort", "POST"), + (f"/session/{job.session_id}", "DELETE")): + try: + _opencode_json_request(job.base_url, path, method=method, + timeout=_LOOP_STOP_TIMEOUT_SECONDS) + except Exception as exc: + # An in-flight turn that refuses to abort still stops + # mattering: the job is gone from the registry, so its + # result is dropped and no further tick is scheduled. + logger.debug("loop %s: %s %s failed: %s", job_id, method, path, exc) return {"ok": True} def update(self, job_id: str, prompt: Optional[str] = None, interval_seconds: Optional[float] = None) -> dict: diff --git a/weightslab/utils/computational_graph.py b/weightslab/utils/computational_graph.py index 34d1699f..65f5b793 100644 --- a/weightslab/utils/computational_graph.py +++ b/weightslab/utils/computational_graph.py @@ -1,3 +1,4 @@ +import inspect import tempfile import torch.nn as nn import torch as th @@ -40,6 +41,12 @@ def _export_model_to_onnx_temp( """ onnx_file_path = tempfile.mkstemp(suffix='.onnx')[1] + # The classic (TorchScript) exporter, explicitly. Since torch 2.9 the default + # is the dynamo exporter, which ignores training=/do_constant_folding= and + # optimizes the graph -- a BatchNorm folded into the conv before it -- so the + # layers the dependency walk reads were gone: a grouped conv's constraint no + # longer reached its BN. `dynamo` itself only exists from torch 2.5. + legacy = {"dynamo": False} if "dynamo" in inspect.signature(th.onnx.export).parameters else {} try: th.onnx.export( model, @@ -51,7 +58,8 @@ def _export_model_to_onnx_temp( input_names=['input'], training=_C_onnx.TrainingMode.PRESERVE, output_names=['output'], - verbose=False + verbose=False, + **legacy, ) logger.info(f"Model exported to {onnx_file_path} (opset {opset_version})") except Exception as e: diff --git a/weightslab/utils/logs.py b/weightslab/utils/logs.py index c62ad66a..ef270a15 100644 --- a/weightslab/utils/logs.py +++ b/weightslab/utils/logs.py @@ -211,6 +211,32 @@ def _print_log_location(): pass +class _SessionFileHandler(logging.FileHandler): + """The session log's handler: logging must never be what fails a program. + + ``FileHandler.emit`` reopens a closed stream *outside* its own try/except, + so once the log's directory was gone -- a temporary ``root_log_dir`` the + log had been moved into, then deleted -- every later log call anywhere in + the process raised ``FileNotFoundError``, failing code that has nothing to + do with logging (tornado creating an event loop, in the release tests). + The directory is recreated when it can be; any other failure goes through + ``handleError``, as a failed write already does. + """ + + def _open(self): + try: + return super()._open() + except FileNotFoundError: + os.makedirs(os.path.dirname(self.baseFilename) or ".", exist_ok=True) + return super()._open() + + def emit(self, record): + try: + super().emit(record) + except Exception: + self.handleError(record) + + def _make_file_handler(path: str) -> logging.FileHandler: """Open the session log file. @@ -224,7 +250,7 @@ def _make_file_handler(path: str) -> logging.FileHandler: equivalent to truncating for a fresh file, and strictly better if two processes ever land on the same name. """ - handler = logging.FileHandler(path, mode='a', encoding='utf-8') + handler = _SessionFileHandler(path, mode='a', encoding='utf-8') handler.setLevel(_FILE_LEVEL) handler.setFormatter(logging.Formatter(FORMAT, datefmt=DATE_FORMAT)) return handler