From d00d97c9fde78c943ba8e1346738a1ff81a3eb0a Mon Sep 17 00:00:00 2001 From: Kevin Heye Date: Thu, 20 Aug 2026 12:16:13 +0200 Subject: [PATCH 01/29] Add AGENTS.md and re-sync the shared CONTRIBUTING block The team-wide part of CONTRIBUTING.md is a copy of templates/CONTRIBUTING.md in the data_team repository, kept in sync by the sync-style-guide skill. Four repos held byte-identical copies and classification_node had already drifted, so the copies were re-synced from one corrected template and the marker heading convention the skill expects is now documented in the vault. AGENTS.md gives an agent the project structure, the commands that actually work here and the traps specific to this repository. It points at CONTRIBUTING.md for the coding standards instead of repeating them a third time. CLAUDE.md imports it via @AGENTS.md, matching nicegui, zoe-app and data_team. Co-authored-by: Claude Opus 5 (1M context) --- AGENTS.md | 56 +++++++++++++++++++++++++++++++++++++++++ CLAUDE.md | 1 + CONTRIBUTING.md | 66 ++++++++++++++++++++++++++++++++++++------------- 3 files changed, 106 insertions(+), 17 deletions(-) create mode 100644 AGENTS.md create mode 100644 CLAUDE.md diff --git a/AGENTS.md b/AGENTS.md new file mode 100644 index 00000000..1762605a --- /dev/null +++ b/AGENTS.md @@ -0,0 +1,56 @@ +# AI Agent Guidelines for the Learning Loop Node Library + +`learning_loop_node` is the **public Python library** every node uses to talk to the +[Learning Loop](https://learning-loop.ai). Four node types build on it: Trainer, Detector, +Annotator and Converter. It is published to PyPI, so its public surface is an API that +`loop`, `dfine_node`, `yolov5_node` and `classification_node` depend on. + +For coding standards see [CONTRIBUTING.md](CONTRIBUTING.md). [README.md](README.md) documents the +environment variables, the node types and how to write a node against them. + +## Layout + +- `learning_loop_node/` — the library itself: `node.py` and `rest.py` for the shared node + machinery, `trainer/`, `detector/`, `annotation/` for the per-type base logic, `data_classes/` + and `enums/` for the wire types, `loop_communication.py` and `data_exchanger.py` for the + loop-facing HTTP and socket.io traffic. +- `learning_loop_node/tests/` — `annotator`, `detector`, `trainer` and `general` suites. +- `mock_trainer/`, `mock_detector/`, `mock_annotator/` — reference implementations with their own + tests. They are what `loop`'s CI runs against, so they are also the best template for a new node. +- `demo_segmentation_tool/` — a worked annotator example. + +## Running and testing + +The suites talk to a real Learning Loop instance and read their credentials from a local `.env` +(`LOOP_HOST`, `LOOP_USERNAME`, `LOOP_PASSWORD`). Without a reachable loop they cannot pass — do not +treat their failure as a regression you introduced. + +```bash +./run_tests.sh # all suites +./run_tests.sh # passed to pytest as -k +``` + +There is no `.pre-commit-config.yaml` here and no ruff in the project environment, despite what the +shared Linting section says. Lint with: + +```bash +uvx ruff check . +``` + +A clean tree already reports several hundred ruff findings, so a clean run is not a reachable goal. +Compare the count on the files you touched, before and after. + +`.github/workflows/pytest.yml` runs the suites, `publish.yml` releases to PyPI on a tagged release. + +## Working in this repository + +- **This is a library — declaring a dependency is part of its API.** Before removing or loosening + one, grep the consumers (`../loop`, `../dfine_node`, `../yolov5_node`, + `../classification_node`) for the package: a consumer that imports it without declaring it + inherits it from here and breaks when it goes away. +- **Renaming or reshaping anything exported** breaks those four repositories. Say so in the pull + request and check whether a companion change is needed there. +- `../loop` checks this repository out as its `nodes` symlink, so a local change is visible to a + local loop immediately — but only a released version reaches CI and production. +- Bump `version` in `pyproject.toml` for a release; the trainer nodes pin the library version in + their image tags (`A.B.C-nlvX.Y.Z`). diff --git a/CLAUDE.md b/CLAUDE.md new file mode 100644 index 00000000..43c994c2 --- /dev/null +++ b/CLAUDE.md @@ -0,0 +1 @@ +@AGENTS.md diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index 5e7f64ed..4ccaa835 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -1,21 +1,35 @@ +# Contributing to Learning Loop Node + ## Data Team Standards ## Linting -We use ruff and [pre-commit](https://github.com/pre-commit/pre-commit) to make sure the coding style is enforced. -You first need to install pre-commit and the corresponding git commit hooks by running the following commands: +We use [ruff](https://docs.astral.sh/ruff/) to enforce the coding style. +Its rules live in the `pyproject.toml` of the respective (sub-)project. + +Repositories that ship a `.pre-commit-config.yaml` run ruff through +[pre-commit](https://github.com/pre-commit/pre-commit). +Install the git hooks once: ```bash python3 -m pip install pre-commit pre-commit install ``` -After that you can make sure your code satisfies the coding style by running the following command: +After that you can check the whole repository with: ```bash pre-commit run --all-files ``` +Repositories without a `.pre-commit-config.yaml` run ruff directly: + +```bash +uvx ruff check . +``` + +The repository-specific part of this file names the exact command wherever it differs. + ## Style Guide ### 1. Single or double quotes @@ -26,8 +40,8 @@ Double quotes may also be used for strings if the string contains a single quote ### 2. String formatting We use f-strings if possible (due to their readability). -Logs shell use the lazy formatting provided by the logging (performance optimization). -When diverging from above rules, please provide a short comment with `# NOTE: ...` to explaining the reason. +Logs shall use the lazy formatting provided by logging (performance optimization). +When diverging from above rules, please provide a short comment with `# NOTE: ...` to explain the reason. ### 3. Line continuation @@ -72,7 +86,7 @@ def parse_config(path: Path) -> Config | None: ... ``` -## 6. Ordering: Important things first +### 6. Ordering: Important things first We want to have main classes and functions at the top of the file, while helper functions and classes should be placed below. This allows to quickly understand the main purpose of the file without having to scroll through a lot of code. @@ -101,23 +115,23 @@ class ReportGenerator: ... ``` -## 7. Docstrings and type hints +### 7. Docstrings and type hints We use the reStructuredText (reST) or Sphinx-style docstring format. We don't declare types in the docstring but use type hints. While type hints are mandatory, parameter descriptions and docstrings should only be used if the function and parameter names are not self-explanatory. -Examples are listetd below: +Examples are listed below: -### One-line Docstring without parameters +#### One-line Docstring without parameters ```python def greet() -> None: """Print a friendly greeting.""" - print("Hello!") + print('Hello!') ``` -### One-line Docstring with parameters +#### One-line Docstring with parameters ```python def square(x: int | float) -> int | float: @@ -129,7 +143,7 @@ def square(x: int | float) -> int | float: return x * x ``` -### Multi-line Docstring without parameters +#### Multi-line Docstring without parameters ```python def get_timestamp() -> str: @@ -140,10 +154,10 @@ def get_timestamp() -> str: as "YYYY-MM-DD HH:MM:SS". It can be useful for logging or displaying time information in user interfaces. """ - return datetime.now().strftime("%Y-%m-%d %H:%M:%S") + return datetime.now().strftime('%Y-%m-%d %H:%M:%S') ``` -### Multi-line Docstring with parameters +#### Multi-line Docstring with parameters ```python def save_data(data: str | bytes, filename: str, overwrite: bool = False) -> None: @@ -157,12 +171,30 @@ def save_data(data: str | bytes, filename: str, overwrite: bool = False) -> None :param data: The content to be saved. Can be text or bytes. :param filename: Path to the destination file. :param overwrite: Whether to overwrite an existing file. - :return: True if the data was saved successfully, False otherwise. :raises FileExistsError: If the file exists and overwrite is False. """ if os.path.exists(filename) and not overwrite: - raise FileExistsError(f"{filename} already exists.") - mode = "wb" if isinstance(data, bytes) else "w" + raise FileExistsError(f'{filename} already exists.') + mode = 'wb' if isinstance(data, bytes) else 'w' with open(filename, mode) as f: f.write(data) ``` + +### 8. Relative imports within a package + +Inside a package we import other modules of the same package with relative imports, not with the package name. +This way the package can be renamed or moved without touching its imports. +We use the shortest path: siblings with a single dot, no going up and back down (`from ..ui.numpad import ...` in a module that lives in `ui/` itself). +Code outside the package — scripts, tests — has no choice and uses absolute imports. + +```python +# preferred (in app/ui/scan_view.py) +from .. import config +from ..system import System +from .numpad import Numpad + +# avoid +from app import config +from app.system import System +from app.ui.numpad import Numpad +``` From 2456d631d065b646b69c9f0c71f7cad9f6c166a9 Mon Sep 17 00:00:00 2001 From: Kevin Heye Date: Thu, 20 Aug 2026 12:40:14 +0200 Subject: [PATCH 02/29] Improve AGENTS.md --- AGENTS.md | 52 ++++++++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 52 insertions(+) diff --git a/AGENTS.md b/AGENTS.md index 1762605a..54866a31 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -19,6 +19,43 @@ environment variables, the node types and how to write a node against them. tests. They are what `loop`'s CI runs against, so they are also the best template for a new node. - `demo_segmentation_tool/` — a worked annotator example. +## Architecture + +Every node is a `FastAPI` subclass (`node.py`): its lifespan connects to the loop and starts a +`repeat_loop` that calls the subclass' `on_repeat` every `repeat_loop_cycle_sec` (5 s) — that loop, +not an event handler, is what drives status reporting, model updates and training continuation. +Subclasses implement `on_startup`, `on_shutdown`, `on_repeat` and `register_sio_events`. + +Two channels lead to the loop and both are needed: `LoopCommunicator` (httpx, login cookies, retry +on 401/429) for the REST API, and a socket.io *client* for status updates and loop-issued commands. +`DataExchanger` sits on top of the communicator and moves images and model zips. + +- **Trainer** — `TrainerLogicGeneric._training_loop` is a state machine over `TrainerState` + (`enums/trainer.py`): download data → download base model → train → sync confusion matrix → + upload model → detect → upload detections → cleanup. `_perform_state` wraps each step: an + ordinary exception records the error and rewinds to the previous state (retried on the next + cycle), a `CriticalError` jumps to `ReadyForCleanup`. Every transition is persisted through + `LastTrainingIO`, so `try_continue_run_if_incomplete` resumes an interrupted training after a + restart. The loop starts a training via the `begin_training` sio event. A concrete trainer + implements `_train`, `_do_detections`, `_get_new_best_training_state`, `_on_metrics_published`, + `_get_latest_model_files` and `_clear_training_data`; `TrainerLogic` adds an `Executor` for + trainers that shell out to a training process. +- **Detector** — the exception to the pattern: constructed with `needs_login=False, needs_sio=False`, + so it has no sio client to the loop. It *hosts* a socket.io server for its own clients and polls + `/{org}/projects/{project}/deployment/target` over REST in `on_repeat` instead. `_DetectorState` + (`_Initializing` / `_Updating` / `_ActiveDetector`) models the model swap: download to + `models/`, build a `DetectorLogic` through the factory, then swap atomically so the old + model keeps serving until the new one is ready (unless `EXCLUSIVE_MODEL_BUILD` frees VRAM first). + `OperationMode` gates whether updates may happen at all. Detections flow through + `RelevanceFilter`, which writes selected images to the `Outbox` on disk; a separate upload process + drains it. +- **Annotator** — thin: it forwards the loop frontend's `handle_user_input` events into + `AnnotatorLogic` and keeps a per-frontend history. + +All node state lives under `GLOBALS.data_folder` (`DATA_FOLDER`, default `/data`): `uuids.json` +(the node uuid is derived from its name and reused across restarts), `models/` plus the +`current_model` symlink, `outbox/`, and the per-project training folders. + ## Running and testing The suites talk to a real Learning Loop instance and read their credentials from a local `.env` @@ -30,6 +67,21 @@ treat their failure as a regression you introduced. ./run_tests.sh # passed to pytest as -k ``` +Each suite carries its own `pytest.ini` (that is where `asyncio_mode = auto` comes from), so always +run pytest with a path inside one suite — a bare `pytest` from the repository root picks up no +config and the async tests error out: + +```bash +python -m pytest learning_loop_node/tests/trainer -v # one suite +python -m pytest learning_loop_node/tests/trainer/test_errors.py -v # one file +python -m pytest learning_loop_node/tests/trainer -v -k # one test +``` + +An autouse fixture repoints `GLOBALS.data_folder` at `/tmp/learning_loop_lib_data` and wipes it +around every test, so tests never touch `/data`. The `general` suite generates and deletes a real +`zauberzeug/pytest_nodelib_general` project on the loop; the detector suite starts the node in a +forked uvicorn process on `GLOBALS.detector_port`. + There is no `.pre-commit-config.yaml` here and no ruff in the project environment, despite what the shared Linting section says. Lint with: From a004efedf42bfc98c10514f782db77acecdc6386 Mon Sep 17 00:00:00 2001 From: Kevin Heye Date: Thu, 20 Aug 2026 15:48:13 +0200 Subject: [PATCH 03/29] Keep AGENTS.md free of private repository names MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit learning_loop_node is public, but AGENTS.md named dfine_node and classification_node, which are not — neither appears anywhere else in the public tree. The dependency rule they illustrated holds without naming them, and "grep ../dfine_node" is advice no external contributor can follow anyway. Also drop the claim that the shared Linting section contradicts this repo: that section now covers repositories without a .pre-commit-config.yaml explicitly. Co-authored-by: Claude Opus 5 (1M context) --- AGENTS.md | 16 +++++++--------- 1 file changed, 7 insertions(+), 9 deletions(-) diff --git a/AGENTS.md b/AGENTS.md index 54866a31..823eeee6 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -2,8 +2,8 @@ `learning_loop_node` is the **public Python library** every node uses to talk to the [Learning Loop](https://learning-loop.ai). Four node types build on it: Trainer, Detector, -Annotator and Converter. It is published to PyPI, so its public surface is an API that -`loop`, `dfine_node`, `yolov5_node` and `classification_node` depend on. +Annotator and Converter. It is published to PyPI, so its public surface is an API that the +Learning Loop backend and every node repository depends on — `../yolov5_node` is the public one. For coding standards see [CONTRIBUTING.md](CONTRIBUTING.md). [README.md](README.md) documents the environment variables, the node types and how to write a node against them. @@ -82,8 +82,7 @@ around every test, so tests never touch `/data`. The `general` suite generates a `zauberzeug/pytest_nodelib_general` project on the loop; the detector suite starts the node in a forked uvicorn process on `GLOBALS.detector_port`. -There is no `.pre-commit-config.yaml` here and no ruff in the project environment, despite what the -shared Linting section says. Lint with: +There is no `.pre-commit-config.yaml` here and no ruff in the project environment. Lint with: ```bash uvx ruff check . @@ -97,11 +96,10 @@ Compare the count on the files you touched, before and after. ## Working in this repository - **This is a library — declaring a dependency is part of its API.** Before removing or loosening - one, grep the consumers (`../loop`, `../dfine_node`, `../yolov5_node`, - `../classification_node`) for the package: a consumer that imports it without declaring it - inherits it from here and breaks when it goes away. -- **Renaming or reshaping anything exported** breaks those four repositories. Say so in the pull - request and check whether a companion change is needed there. + one, grep every consuming repository checked out beside this one for the package: a consumer that + imports it without declaring it inherits it from here and breaks when it goes away. +- **Renaming or reshaping anything exported** breaks those repositories. Say so in the pull request + and check whether a companion change is needed there. - `../loop` checks this repository out as its `nodes` symlink, so a local change is visible to a local loop immediately — but only a released version reaches CI and production. - Bump `version` in `pyproject.toml` for a release; the trainer nodes pin the library version in From 7747813a0387db9a38d083aee158727d23da8d06 Mon Sep 17 00:00:00 2001 From: Kevin Heye Date: Thu, 20 Aug 2026 16:47:01 +0200 Subject: [PATCH 04/29] Share detector postprocessing across nodes Every detector node does the same three things with a model's raw output: drop low-confidence predictions, suppress overlapping boxes, and turn what survives into the loop's dataclasses. None of it depends on the model, yet each node repository had grown its own copy, and the copies had drifted: clip_point was byte-identical in three places, while clip_box existed in the same three plus a fourth that used a centre-based convention under the same name. Move that code here as detector/postprocess.py and detector/geometry.py, and give the centre-based variant its own name (clip_box_centered) so the two conventions can no longer be confused. The library also carried two containers for the same detections: a detector reports ImageMetadata, a trainer's auto-detection pass reports Detections. to_image_metadata and to_detections now build both from one routine, so the two paths clip and filter identically instead of diverging. The scalar arguments of non_max_suppression and post_process are keyword-only: transposing origin_h and origin_w was too easy to do. Add tests/unit, the first suite that runs without a Learning Loop, and gate the credential-dependent jobs on it in CI. Co-Authored-By: Claude Opus 5 (1M context) --- .github/workflows/pytest.yml | 22 ++ AGENTS.md | 28 +- learning_loop_node/detector/geometry.py | 69 +++++ learning_loop_node/detector/postprocess.py | 264 ++++++++++++++++++ learning_loop_node/tests/unit/__init__.py | 0 learning_loop_node/tests/unit/pytest.ini | 10 + .../tests/unit/test_geometry.py | 35 +++ .../tests/unit/test_postprocess.py | 168 +++++++++++ run_tests.sh | 3 + 9 files changed, 594 insertions(+), 5 deletions(-) create mode 100644 learning_loop_node/detector/geometry.py create mode 100644 learning_loop_node/detector/postprocess.py create mode 100644 learning_loop_node/tests/unit/__init__.py create mode 100644 learning_loop_node/tests/unit/pytest.ini create mode 100644 learning_loop_node/tests/unit/test_geometry.py create mode 100644 learning_loop_node/tests/unit/test_postprocess.py diff --git a/.github/workflows/pytest.yml b/.github/workflows/pytest.yml index 3e456ca2..e334d0b4 100644 --- a/.github/workflows/pytest.yml +++ b/.github/workflows/pytest.yml @@ -3,7 +3,28 @@ name: Run Tests on: [push] jobs: + unit: + runs-on: ubuntu-latest + timeout-minutes: 10 + steps: + - uses: actions/checkout@v4 + - name: set up Python + uses: actions/setup-python@v5 + with: + python-version: "3.10" + - name: set up uv + uses: astral-sh/setup-uv@v3 + - name: install dependencies + run: | + uv sync --extra dev --frozen + - name: test_unit + # no Learning Loop and no secrets needed, so this suite gates the slow ones + run: | + uv run python -m pytest learning_loop_node/tests/unit -v + pytest_3_10: + needs: + - unit runs-on: ubuntu-latest timeout-minutes: 30 strategy: @@ -91,6 +112,7 @@ jobs: slack: needs: + - unit - pytest_3_10 - pytest_3_13 if: always() # also execute when pytest fails diff --git a/AGENTS.md b/AGENTS.md index 823eeee6..cf26c812 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -14,7 +14,9 @@ environment variables, the node types and how to write a node against them. machinery, `trainer/`, `detector/`, `annotation/` for the per-type base logic, `data_classes/` and `enums/` for the wire types, `loop_communication.py` and `data_exchanger.py` for the loop-facing HTTP and socket.io traffic. -- `learning_loop_node/tests/` — `annotator`, `detector`, `trainer` and `general` suites. +- `learning_loop_node/tests/` — `unit`, `annotator`, `detector`, `trainer` and `general` suites. + Only `unit` runs without a Learning Loop; it covers the pure helpers such as + `detector/postprocess.py` and `detector/geometry.py`. - `mock_trainer/`, `mock_detector/`, `mock_annotator/` — reference implementations with their own tests. They are what `loop`'s CI runs against, so they are also the best template for a new node. - `demo_segmentation_tool/` — a worked annotator example. @@ -52,15 +54,30 @@ on 401/429) for the REST API, and a socket.io *client* for status updates and lo - **Annotator** — thin: it forwards the loop frontend's `handle_user_input` events into `AnnotatorLogic` and keeps a per-frontend history. +`detector/postprocess.py` and `detector/geometry.py` hold the parts of a detector that do *not* +depend on the model: confidence filtering, per-class NMS, box/point clipping, and turning +predictions into the loop's dataclasses. A node should import them rather than write its own — +every node repository had grown its own drifting copy, which is why they live here. Note the two +containers: `to_image_metadata` builds what a **detector** node reports, `to_detections` what a +**trainer**'s auto-detection pass reports. Both go through one routine, so the two paths cannot +drift apart again. + All node state lives under `GLOBALS.data_folder` (`DATA_FOLDER`, default `/data`): `uuids.json` (the node uuid is derived from its name and reused across restarts), `models/` plus the `current_model` symlink, `outbox/`, and the per-project training folders. ## Running and testing -The suites talk to a real Learning Loop instance and read their credentials from a local `.env` -(`LOOP_HOST`, `LOOP_USERNAME`, `LOOP_PASSWORD`). Without a reachable loop they cannot pass — do not -treat their failure as a regression you introduced. +The `unit` suite is self-contained — no Learning Loop, no credentials, no network — so it is the +one to run while iterating, and the one CI gates the others on: + +```bash +python -m pytest learning_loop_node/tests/unit -v +``` + +Every other suite talks to a real Learning Loop instance and reads its credentials from a local +`.env` (`LOOP_HOST`, `LOOP_USERNAME`, `LOOP_PASSWORD`). Without a reachable loop those cannot pass — +do not treat their failure as a regression you introduced. ```bash ./run_tests.sh # all suites @@ -91,7 +108,8 @@ uvx ruff check . A clean tree already reports several hundred ruff findings, so a clean run is not a reachable goal. Compare the count on the files you touched, before and after. -`.github/workflows/pytest.yml` runs the suites, `publish.yml` releases to PyPI on a tagged release. +`.github/workflows/pytest.yml` runs the suites (the `unit` job first, without secrets), +`publish.yml` releases to PyPI on a tagged release. ## Working in this repository diff --git a/learning_loop_node/detector/geometry.py b/learning_loop_node/detector/geometry.py new file mode 100644 index 00000000..b17ed33b --- /dev/null +++ b/learning_loop_node/detector/geometry.py @@ -0,0 +1,69 @@ +"""Box and point clipping shared by every detector node. + +The loop stores a box as its top-left corner plus a size, so :func:`clip_box` is the form a +node needs when it hands detections to the loop. Model outputs are not always in that form — +:func:`clip_box_centered` keeps the centre-based convention explicit instead of letting two +incompatible functions share one name, which is how the same helper ended up meaning two +different things in different node repositories. +""" + + +def clip_box( + *, + x1: float, + y1: float, + width: float, + height: float, + img_width: int, + img_height: int, +) -> tuple[int, int, int, int]: + """Clip a top-left-anchored box to the image bounds. + + :param x1: Left edge of the box. + :param y1: Top edge of the box. + :return: The clipped ``(x1, y1, width, height)`` as ints; the size is never negative. + """ + x2 = x1 + width + y2 = y1 + height + + clipped_x1 = round(max(0.0, x1)) + clipped_y1 = round(max(0.0, y1)) + clipped_x2 = round(min(float(img_width), x2)) + clipped_y2 = round(min(float(img_height), y2)) + + clipped_width = max(clipped_x2 - clipped_x1, 0) + clipped_height = max(clipped_y2 - clipped_y1, 0) + + return clipped_x1, clipped_y1, clipped_width, clipped_height + + +def clip_box_centered( + *, + x: float, + y: float, + width: float, + height: float, + img_width: int, + img_height: int, +) -> tuple[float, float, float, float]: + """Clip a centre-anchored box to the image bounds, keeping it centre-anchored. + + Clipping moves the centre, because only the part of the box inside the image survives. + + :param x: Horizontal centre of the box. + :param y: Vertical centre of the box. + :return: The clipped ``(x, y, width, height)``, still centre-anchored. + """ + left = max(0.0, x - 0.5 * width) + top = max(0.0, y - 0.5 * height) + right = min(float(img_width), x + 0.5 * width) + bottom = min(float(img_height), y + 0.5 * height) + + return 0.5 * (left + right), 0.5 * (top + bottom), right - left, bottom - top + + +def clip_point(x: float, y: float, img_width: int, img_height: int) -> tuple[float, float]: + """Clamp a point into the image bounds.""" + x = min(max(0, x), img_width) + y = min(max(0, y), img_height) + return x, y diff --git a/learning_loop_node/detector/postprocess.py b/learning_loop_node/detector/postprocess.py new file mode 100644 index 00000000..ae47fbcc --- /dev/null +++ b/learning_loop_node/detector/postprocess.py @@ -0,0 +1,264 @@ +"""Model-agnostic detection postprocessing. + +Every detection node ends up doing the same three things with a model's raw output: drop +low-confidence predictions, suppress overlapping boxes, and turn what survives into the +loop's detection dataclasses. None of that depends on the model, so it lives here rather +than being re-derived — and re-diverging — in each node repository. + +Two containers carry the same detections in this library: a detector node reports +:class:`~learning_loop_node.data_classes.image_metadata.ImageMetadata`, while a trainer's +auto-detection pass reports :class:`~learning_loop_node.data_classes.detections.Detections`. +:func:`to_image_metadata` and :func:`to_detections` build them from the same routine, so both +paths clip and filter identically. +""" + +import logging +from collections import namedtuple + +import numpy as np + +from ..data_classes import ( + BoxDetection, + Category, + Detections, + ImageMetadata, + ModelInformation, + PointDetection, +) +from ..enums import CategoryType +from .geometry import clip_box, clip_point + +MIN_BOX_SIZE: int = 2 +"""Boxes this small are dropped: they carry no usable information and clutter the loop.""" + +Detection = namedtuple('Detection', 'x y w h category probability') +"""One surviving prediction. ``x``/``y`` are the top-left corner, ``category`` is an index +into :attr:`ModelInformation.categories`.""" + + +def category_by_index(model_information: ModelInformation, index: int) -> Category: + """Resolve the category a model's class index refers to. + + Models emit class indices, and the order of ``model_information.categories`` is what + gives them meaning — so an out-of-range index is a model/metadata mismatch, not a + detection to skip quietly. + + :raises ValueError: If the index is outside the model's category list. + """ + categories = model_information.categories + if not 0 <= index < len(categories): + raise ValueError( + f'category index {index} is out of range for a model with {len(categories)} categories') + return categories[index] + + +def category_by_name(model_information: ModelInformation, name: str) -> Category: + """Resolve a category by name, for models whose outputs are named rather than indexed. + + :raises ValueError: If no category of that name exists. + """ + for category in model_information.categories: + if category.name == name: + return category + known = ', '.join(category.name for category in model_information.categories) + raise ValueError(f'unknown category name {name!r}; the model knows: {known}') + + +def bbox_iou( + box1: np.ndarray, + box2: np.ndarray, +) -> np.ndarray: + """Compute IoU between box1 (1x4) and box2 (Nx4), both in x1y1x2y2 format.""" + b1_x1, b1_y1, b1_x2, b1_y2 = box1[:, 0], box1[:, 1], box1[:, 2], box1[:, 3] + b2_x1, b2_y1, b2_x2, b2_y2 = box2[:, 0], box2[:, 1], box2[:, 2], box2[:, 3] + + inter_x1 = np.maximum(b1_x1, b2_x1) + inter_y1 = np.maximum(b1_y1, b2_y1) + inter_x2 = np.minimum(b1_x2, b2_x2) + inter_y2 = np.minimum(b1_y2, b2_y2) + + inter_area = np.clip(inter_x2 - inter_x1 + 1, 0, None) * np.clip(inter_y2 - inter_y1 + 1, 0, None) + b1_area = (b1_x2 - b1_x1 + 1) * (b1_y2 - b1_y1 + 1) + b2_area = (b2_x2 - b2_x1 + 1) * (b2_y2 - b2_y1 + 1) + + return inter_area / (b1_area + b2_area - inter_area + 1e-16) + + +def non_max_suppression( + boxes: np.ndarray, + scores: np.ndarray, + classes: np.ndarray, + *, + iou_threshold: float, + origin_h: int, + origin_w: int, +) -> tuple[np.ndarray, np.ndarray, np.ndarray]: + """Clip to image bounds, sort by score descending, apply per-class NMS. + + :return: ``(boxes, scores, classes)`` — filtered arrays in the same order. + """ + boxes = boxes.copy() + boxes[:, 0] = np.clip(boxes[:, 0], 0, origin_w - 1) + boxes[:, 2] = np.clip(boxes[:, 2], 0, origin_w - 1) + boxes[:, 1] = np.clip(boxes[:, 1], 0, origin_h - 1) + boxes[:, 3] = np.clip(boxes[:, 3], 0, origin_h - 1) + + order = np.argsort(-scores) + boxes = boxes[order] + scores = scores[order] + classes = classes[order] + + keep_indices: list[int] = [] + for cls in np.unique(classes): + cls_mask = classes == cls + cls_indices = np.where(cls_mask)[0] + cls_boxes = boxes[cls_mask] + + while len(cls_indices) > 0: + keep_indices.append(cls_indices[0]) + if len(cls_indices) == 1: + break + ious = bbox_iou(cls_boxes[0:1], cls_boxes[1:]) + keep_mask = ious.flatten() <= iou_threshold + cls_indices = cls_indices[1:][keep_mask] + cls_boxes = cls_boxes[1:][keep_mask] + + keep_indices = sorted(keep_indices) + return boxes[keep_indices], scores[keep_indices], classes[keep_indices] + + +def post_process( + boxes: np.ndarray, + scores: np.ndarray, + classes: np.ndarray, + *, + conf_threshold: float, + iou_threshold: float, + origin_h: int, + origin_w: int, +) -> list[Detection]: + """Filter by confidence, run NMS, return a :class:`Detection` list in x/y/w/h form.""" + mask = scores > conf_threshold + boxes = boxes[mask].copy() + scores = scores[mask] + classes = classes[mask] + + if len(scores) == 0: + return [] + + boxes, scores, classes = non_max_suppression( + boxes, scores, classes, + iou_threshold=iou_threshold, origin_h=origin_h, origin_w=origin_w) + + result = [] + for j, box in enumerate(boxes): + x1, y1, x2, y2 = box + w = x2 - x1 + h = y2 - y1 + result.append(Detection(int(x1), int(y1), int(w), int(h), + int(classes[j]), round(float(scores[j]), 2))) + return result + + +def detections_from_xyxy( + *, + labels: list[float], + boxes: list[list[float]], + scores: list[float], +) -> list[Detection]: + """Convert already-suppressed model output into :class:`Detection` values. + + For nodes whose model (or a torch/ONNX op) has done the suppression already, so only the + coordinate conversion is left. Corners are rounded rather than truncated, which is half a + pixel more faithful than :func:`post_process` — that one keeps truncating so its output + stays bit-identical to what detectors reported before this module existed. + """ + result = [] + for label, box, score in zip(labels, boxes, scores, strict=True): + x1, y1, x2, y2 = (round(value) for value in box) + result.append(Detection(x1, y1, x2 - x1, y2 - y1, int(label), score)) + return result + + +def _append_detections( + target: ImageMetadata | Detections, + detections: list[Detection], + model_information: ModelInformation, + im_height: int, + im_width: int, +) -> None: + """Resolve each detection's category and append it to ``target``, clipped to the image.""" + skipped_detections = [] + + for detection in detections: + x, y, w, h, category_idx, probability = detection + category = category_by_index(model_information, category_idx) + if w <= MIN_BOX_SIZE or h <= MIN_BOX_SIZE: + skipped_detections.append((category.name, detection)) + continue + if category.type == CategoryType.Box: + clipped_x1, clipped_y1, clipped_w, clipped_h = clip_box( + x1=x, + y1=y, + width=w, + height=h, + img_width=im_width, + img_height=im_height, + ) + target.box_detections.append( + BoxDetection( + category_name=category.name, + x=clipped_x1, + y=clipped_y1, + width=clipped_w, + height=clipped_h, + category_id=category.id, + model_name=model_information.version, + confidence=probability, + ) + ) + elif category.type == CategoryType.Point: + cx, cy = x + w / 2, y + h / 2 + cx, cy = clip_point(cx, cy, im_width, im_height) + target.point_detections.append( + PointDetection( + category_name=category.name, + x=cx, + y=cy, + category_id=category.id, + model_name=model_information.version, + confidence=probability, + ) + ) + else: + logging.warning('Unsupported category type %s for category %s', category.type, category.name) + + if skipped_detections: + log_msg = '\n'.join([str(d) for d in skipped_detections]) + logging.warning('Removed %d small detections from result: \n%s', len(skipped_detections), log_msg) + + +def to_image_metadata( + detections: list[Detection], + model_information: ModelInformation, + im_height: int, + im_width: int, +) -> ImageMetadata: + """Build the container a *detector* node reports from a list of detections.""" + image_metadata = ImageMetadata() + _append_detections(image_metadata, detections, model_information, im_height, im_width) + return image_metadata + + +def to_detections( + detections: list[Detection], + model_information: ModelInformation, + im_height: int, + im_width: int, + *, + image_id: str | None = None, +) -> Detections: + """Build the container a *trainer*'s auto-detection pass reports.""" + result = Detections(image_id=image_id) + _append_detections(result, detections, model_information, im_height, im_width) + return result diff --git a/learning_loop_node/tests/unit/__init__.py b/learning_loop_node/tests/unit/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/learning_loop_node/tests/unit/pytest.ini b/learning_loop_node/tests/unit/pytest.ini new file mode 100644 index 00000000..2207d1ac --- /dev/null +++ b/learning_loop_node/tests/unit/pytest.ini @@ -0,0 +1,10 @@ +[pytest] +python_files = test_*.py +asyncio_mode = auto + +cache_dir = /tmp/pytest_cache + +# for debbuging tests: +; log_cli_level = INFO +; log_cli_format = %(asctime)s [%(levelname)8s] %(message)s (%(filename)s:%(lineno)s) +; log_cli_date_format=%Y-%m-%d %H:%M:%S \ No newline at end of file diff --git a/learning_loop_node/tests/unit/test_geometry.py b/learning_loop_node/tests/unit/test_geometry.py new file mode 100644 index 00000000..f056f79d --- /dev/null +++ b/learning_loop_node/tests/unit/test_geometry.py @@ -0,0 +1,35 @@ +from ...detector.geometry import clip_box, clip_box_centered, clip_point + + +def test_box_inside_the_image_is_unchanged(): + assert clip_box(x1=10, y1=20, width=30, height=40, img_width=100, img_height=100) == (10, 20, 30, 40) + + +def test_box_is_clipped_to_the_image_bounds(): + assert clip_box(x1=-20, y1=-20, width=60, height=60, img_width=100, img_height=100) == (0, 0, 40, 40) + assert clip_box(x1=80, y1=80, width=40, height=40, img_width=100, img_height=100) == (80, 80, 20, 20) + + +def test_box_fully_outside_the_image_collapses_to_zero_size(): + # The corner is only clamped at the lower bound, so it stays at 200 — the zero size is + # what marks the box as empty. Detector output cannot reach here, because + # non_max_suppression already clips every box into the image. + assert clip_box(x1=200, y1=200, width=10, height=10, img_width=100, img_height=100) == (200, 200, 0, 0) + + +def test_box_corners_are_rounded(): + assert clip_box(x1=10.4, y1=10.6, width=20.0, height=20.0, img_width=100, img_height=100) == (10, 11, 20, 20) + + +def test_centered_box_keeps_its_centre_when_it_fits(): + assert clip_box_centered(x=50, y=50, width=20, height=20, img_width=100, img_height=100) == (50, 50, 20, 20) + + +def test_clipping_a_centered_box_moves_its_centre(): + # only the right half of the box is inside the image, so the centre moves right + assert clip_box_centered(x=0, y=50, width=20, height=20, img_width=100, img_height=100) == (5, 50, 10, 20) + + +def test_point_is_clamped_into_the_image(): + assert clip_point(50, 50, 100, 100) == (50, 50) + assert clip_point(-10, 150, 100, 100) == (0, 100) diff --git a/learning_loop_node/tests/unit/test_postprocess.py b/learning_loop_node/tests/unit/test_postprocess.py new file mode 100644 index 00000000..14dbcd6a --- /dev/null +++ b/learning_loop_node/tests/unit/test_postprocess.py @@ -0,0 +1,168 @@ +import numpy as np +import pytest + +from ...data_classes import Category, ModelInformation +from ...detector.postprocess import ( + Detection, + bbox_iou, + category_by_index, + category_by_name, + detections_from_xyxy, + non_max_suppression, + post_process, + to_detections, + to_image_metadata, +) +from ...enums import CategoryType + +BOX = Category(id='uuid-box', name='car', type=CategoryType.Box) +POINT = Category(id='uuid-point', name='weed', type=CategoryType.Point) + + +def model_information(*categories: Category) -> ModelInformation: + return ModelInformation(id='model-uuid', host='localhost', organization='zauberzeug', + project='pytest', version='1.2', categories=list(categories or (BOX, POINT))) + + +# ---------------------------------------------------------------- category resolution + +def test_category_is_resolved_by_index(): + assert category_by_index(model_information(), 1) is POINT + + +@pytest.mark.parametrize('index', [-1, 2, 99]) +def test_an_index_outside_the_model_categories_is_an_error(index: int): + # a mismatch between model and metadata must not be silently skipped + with pytest.raises(ValueError, match='out of range'): + category_by_index(model_information(), index) + + +def test_category_is_resolved_by_name(): + assert category_by_name(model_information(), 'weed') is POINT + + +def test_an_unknown_category_name_lists_the_known_ones(): + with pytest.raises(ValueError, match='car, weed'): + category_by_name(model_information(), 'tractor') + + +# ---------------------------------------------------------------- iou and suppression + +def test_identical_boxes_have_an_iou_of_one(): + box = np.array([[0, 0, 10, 10]], dtype=np.float32) + assert bbox_iou(box, box)[0] == pytest.approx(1.0) + + +def test_disjoint_boxes_have_an_iou_of_zero(): + a = np.array([[0, 0, 10, 10]], dtype=np.float32) + b = np.array([[100, 100, 110, 110]], dtype=np.float32) + assert bbox_iou(a, b)[0] == pytest.approx(0.0) + + +def test_overlapping_boxes_of_one_class_are_suppressed_keeping_the_best(): + boxes = np.array([[10, 10, 60, 60], [12, 12, 62, 62]], dtype=np.float32) + scores = np.array([0.9, 0.8], dtype=np.float32) + kept_boxes, kept_scores, _ = non_max_suppression( + boxes, scores, np.array([0, 0]), iou_threshold=0.45, origin_h=200, origin_w=200) + assert len(kept_boxes) == 1 + assert kept_scores[0] == pytest.approx(0.9) + + +def test_overlapping_boxes_of_different_classes_both_survive(): + boxes = np.array([[10, 10, 60, 60], [12, 12, 62, 62]], dtype=np.float32) + kept_boxes, _, _ = non_max_suppression( + boxes, np.array([0.9, 0.8], dtype=np.float32), np.array([0, 1]), + iou_threshold=0.45, origin_h=200, origin_w=200) + assert len(kept_boxes) == 2 + + +def test_suppression_clips_boxes_into_the_image(): + boxes = np.array([[-10, -10, 300, 300]], dtype=np.float32) + kept_boxes, _, _ = non_max_suppression( + boxes, np.array([0.9], dtype=np.float32), np.array([0]), + iou_threshold=0.45, origin_h=100, origin_w=100) + assert list(kept_boxes[0]) == [0, 0, 99, 99] + + +# ---------------------------------------------------------------- post_process + +def test_post_process_drops_predictions_below_the_confidence_threshold(): + boxes = np.array([[10, 10, 60, 60], [100, 100, 150, 150]], dtype=np.float32) + result = post_process(boxes, np.array([0.9, 0.1], dtype=np.float32), np.array([0, 0]), + conf_threshold=0.5, iou_threshold=0.45, origin_h=200, origin_w=200) + assert result == [Detection(10, 10, 50, 50, 0, 0.9)] + + +def test_post_process_on_an_empty_prediction_returns_nothing(): + empty_boxes = np.zeros((0, 4), dtype=np.float32) + assert post_process(empty_boxes, np.zeros(0, dtype=np.float32), np.zeros(0, dtype=int), + conf_threshold=0.5, iou_threshold=0.45, origin_h=10, origin_w=10) == [] + + +def test_already_suppressed_output_is_converted_with_rounded_corners(): + assert detections_from_xyxy(labels=[1.0], boxes=[[10.4, 10.6, 60.4, 60.6]], scores=[0.55]) == \ + [Detection(10, 11, 50, 50, 1, 0.55)] + + +def test_converting_already_suppressed_output_requires_matching_lengths(): + with pytest.raises(ValueError): + detections_from_xyxy(labels=[1.0, 2.0], boxes=[[0.0, 0.0, 1.0, 1.0]], scores=[0.5]) + + +# ---------------------------------------------------------------- building the containers + +def test_a_box_category_becomes_a_box_detection(): + metadata = to_image_metadata([Detection(10, 20, 30, 40, 0, 0.9)], model_information(), 200, 200) + assert len(metadata.point_detections) == 0 + detection = metadata.box_detections[0] + assert (detection.x, detection.y, detection.width, detection.height) == (10, 20, 30, 40) + assert (detection.category_name, detection.category_id) == ('car', 'uuid-box') + assert detection.model_name == '1.2' + assert detection.confidence == pytest.approx(0.9) + + +def test_a_point_category_becomes_the_centre_of_the_box(): + metadata = to_image_metadata([Detection(100, 100, 40, 40, 1, 0.7)], model_information(), 200, 200) + assert len(metadata.box_detections) == 0 + detection = metadata.point_detections[0] + assert (detection.x, detection.y) == (120, 120) + assert detection.category_id == 'uuid-point' + + +def test_detections_are_clipped_to_the_image(): + metadata = to_image_metadata([Detection(-20, -20, 60, 60, 0, 0.5)], model_information(), 200, 200) + detection = metadata.box_detections[0] + assert (detection.x, detection.y, detection.width, detection.height) == (0, 0, 40, 40) + + +@pytest.mark.parametrize('width,height', [(2, 30), (30, 2), (1, 1)]) +def test_boxes_too_small_to_be_useful_are_dropped(width: int, height: int): + metadata = to_image_metadata([Detection(5, 5, width, height, 0, 0.5)], model_information(), 200, 200) + assert len(metadata) == 0 + + +def test_a_category_type_the_node_cannot_report_is_skipped(): + classification = Category(id='uuid-cls', name='ripe', type=CategoryType.Classification) + metadata = to_image_metadata([Detection(10, 10, 30, 30, 0, 0.5)], + model_information(classification), 200, 200) + assert len(metadata) == 0 + + +def test_the_trainer_container_carries_the_image_id(): + result = to_detections([Detection(10, 20, 30, 40, 0, 0.9)], model_information(), 200, 200, + image_id='image-uuid') + assert result.image_id == 'image-uuid' + assert len(result.box_detections) == 1 + + +def test_trainer_and_detector_paths_agree_on_the_same_detections(): + """The whole point of sharing this code: auto-detections and live detections must match.""" + detections = [Detection(-5, -5, 60, 60, 0, 0.9), Detection(100, 100, 40, 40, 1, 0.7), + Detection(5, 5, 1, 1, 0, 0.5)] + metadata = to_image_metadata(detections, model_information(), 200, 200) + result = to_detections(detections, model_information(), 200, 200, image_id='image-uuid') + + assert [(d.x, d.y, d.width, d.height, d.category_id) for d in result.box_detections] == \ + [(d.x, d.y, d.width, d.height, d.category_id) for d in metadata.box_detections] + assert [(d.x, d.y, d.category_id) for d in result.point_detections] == \ + [(d.x, d.y, d.category_id) for d in metadata.point_detections] diff --git a/run_tests.sh b/run_tests.sh index e83d3852..e8ae7b66 100755 --- a/run_tests.sh +++ b/run_tests.sh @@ -6,6 +6,7 @@ set -o allexport; source .env; set +o allexport # Check if argument is provided if [ $# -eq 1 ]; then # Run tests with filter + python -m pytest learning_loop_node/tests/unit -v -s -k "$1" python -m pytest learning_loop_node/tests/annotator -v -s -k "$1" python -m pytest learning_loop_node/tests/detector -v -s -k "$1" python -m pytest learning_loop_node/tests/trainer -v -s -k "$1" @@ -17,6 +18,8 @@ fi # Run the tests +# unit runs first: it is the only suite that needs no Learning Loop +python -m pytest learning_loop_node/tests/unit -v python -m pytest learning_loop_node/tests/annotator -v python -m pytest learning_loop_node/tests/detector -v python -m pytest learning_loop_node/tests/trainer -v From fa8505c1d76170ec20e50aa8e522c1615dfee3b3 Mon Sep 17 00:00:00 2001 From: Kevin Heye Date: Thu, 20 Aug 2026 19:18:53 +0200 Subject: [PATCH 05/29] Share the node entry point boilerplate Every main.py in every node repository repeats the same twenty lines: read UVICORN_RELOAD, build the settings, construct the node, call uvicorn.run. Eight copies across three repositories, and they had drifted - dfine parses settings with configargparse so each one is both a flag and an environment variable, while the others hand-roll os.getenv plus asserts and hardcode the host and port. node_parser and run_node put that in one place. A setting is named after its flag: --conf-threshold reads CONF_THRESHOLD. That is what yolov5_node and classification_node already use, so adopting this renames nothing in the repositories that are actually deployed. --host and --port are the exception, reading NODE_HOST and NODE_PORT. The bare HOST already means the loop's address - every deployment sets it - and a node adopting that would hand it to uvicorn and fail to bind. dfine_node required a DFINE_TRAINER_ / DFINE_DETECTOR_ prefix, so it passes legacy_env_prefix and its prefixed names keep working with a warning naming the replacement. The logging.basicConfig call the trainers carry is not reproduced here - the library configures the root logger on import, so it has been a no-op. Co-Authored-By: Claude Opus 5 (1M context) --- AGENTS.md | 8 ++ learning_loop_node/helpers/entrypoint.py | 75 +++++++++++++++++++ .../tests/unit/test_entrypoint.py | 71 ++++++++++++++++++ pyproject.toml | 1 + uv.lock | 11 +++ 5 files changed, 166 insertions(+) create mode 100644 learning_loop_node/helpers/entrypoint.py create mode 100644 learning_loop_node/tests/unit/test_entrypoint.py diff --git a/AGENTS.md b/AGENTS.md index cf26c812..c7877d5d 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -54,6 +54,14 @@ on 401/429) for the REST API, and a socket.io *client* for status updates and lo - **Annotator** — thin: it forwards the loop frontend's `handle_user_input` events into `AnnotatorLogic` and keeps a per-frontend history. +`helpers/entrypoint.py` holds what every node's `main.py` repeats: `node_parser` builds a +configargparse parser with `--host`/`--port`, `run_node` starts uvicorn. A setting is a flag +*and* an environment variable from one declaration — `--conf-threshold` reads +`CONF_THRESHOLD`. The exception is `--host`/`--port`, which read `NODE_HOST`/`NODE_PORT`: the +bare `HOST` already means the loop's address, and a node adopting it would hand it to uvicorn +and fail to bind. A node that used to require a prefix passes `legacy_env_prefix`, and the +prefixed names keep working with a warning. + `detector/postprocess.py` and `detector/geometry.py` hold the parts of a detector that do *not* depend on the model: confidence filtering, per-class NMS, box/point clipping, and turning predictions into the loop's dataclasses. A node should import them rather than write its own — diff --git a/learning_loop_node/helpers/entrypoint.py b/learning_loop_node/helpers/entrypoint.py new file mode 100644 index 00000000..830c89ec --- /dev/null +++ b/learning_loop_node/helpers/entrypoint.py @@ -0,0 +1,75 @@ +"""The boilerplate every node's ``main.py`` repeats. + +A node entry point always does the same four things: read a handful of settings, build the +logic object, construct the node, and hand it to uvicorn. Only the middle two are the node's +own, so :func:`node_parser` and :func:`run_node` cover the rest. + +Settings come from a flag *or* an environment variable, because a node is configured on the +command line while developing and through the container environment in deployment. One +declaration gives both: ``--conf-threshold`` reads ``CONF_THRESHOLD``. + +``--host`` and ``--port`` are the exception — they read ``NODE_HOST`` and ``NODE_PORT`` rather +than the names their flags imply, because the bare ``HOST`` already means *the loop's address* +and a node adopting it would hand it to uvicorn and fail to bind. +""" + +import logging +import os +from argparse import Action, Namespace + +import configargparse +import uvicorn + + +def node_parser(*, description: str, legacy_env_prefix: str = '') -> configargparse.ArgumentParser: + """Build the parser for a node, pre-loaded with the settings every node has. + + :param legacy_env_prefix: A prefix an earlier version of this node required, e.g. + ``'DFINE_DETECTOR_'``. Prefixed names are still honoured, with a warning, so a + deployment keeps working until it is updated. Leave empty for a node that has always + read unprefixed names. + """ + parser = _NodeArgumentParser(description=description, legacy_env_prefix=legacy_env_prefix) + parser.add_argument('--host', default='0.0.0.0', env_var='NODE_HOST', + help='Host interface to bind to') + parser.add_argument('--port', type=int, default=80, env_var='NODE_PORT', help='Port to bind to') + return parser + + +def run_node(app: str, args: Namespace) -> None: + """Serve the node. + + :param app: Import string of the node object, conventionally ``'main:node'``. + """ + reload = os.getenv('UVICORN_RELOAD', 'FALSE').lower() in ('true', '1') + logging.info('Uvicorn reload is set to: %s', reload) + uvicorn.run(app, host=args.host, port=args.port, lifespan='on', reload=reload) + + +class _NodeArgumentParser(configargparse.ArgumentParser): + """Parser whose every setting is also an environment variable named after its flag.""" + + def __init__(self, *, description: str, legacy_env_prefix: str) -> None: + super().__init__(description=description) + self.legacy_env_prefix = legacy_env_prefix + + def add_argument(self, *args, **kwargs) -> Action: # type: ignore[override] + action = super().add_argument(*args, **kwargs) + if getattr(action, 'env_var', None) is None and action.dest != 'help': + action.env_var = action.dest.upper() + return action + + def parse_args(self, *args, **kwargs) -> Namespace: # type: ignore[override] + self._adopt_legacy_env_vars() + return super().parse_args(*args, **kwargs) + + def _adopt_legacy_env_vars(self) -> None: + if not self.legacy_env_prefix: + return + for action in self._actions: + name = getattr(action, 'env_var', None) + legacy = self.legacy_env_prefix + name if name else None + if not legacy or name in os.environ or legacy not in os.environ: + continue + os.environ[name] = os.environ[legacy] + logging.warning('%s is deprecated and will stop being read; set %s instead', legacy, name) diff --git a/learning_loop_node/tests/unit/test_entrypoint.py b/learning_loop_node/tests/unit/test_entrypoint.py new file mode 100644 index 00000000..a478563b --- /dev/null +++ b/learning_loop_node/tests/unit/test_entrypoint.py @@ -0,0 +1,71 @@ +import pytest + +from ...helpers.entrypoint import node_parser + +MANAGED = ('WEIGHT_TYPE', 'DFINE_DETECTOR_WEIGHT_TYPE', 'HOST', 'NODE_HOST', 'NODE_PORT', 'PORT') + + +@pytest.fixture(autouse=True) +def clean_env(monkeypatch: pytest.MonkeyPatch): + """Every test starts without the variables it is about to set.""" + for name in MANAGED: + monkeypatch.delenv(name, raising=False) + + +def _parser(**kwargs): + parser = node_parser(description='a node', **kwargs) + parser.add_argument('--weight-type', default='FP16') + return parser + + +def test_every_node_gets_a_host_and_a_port(): + args = _parser().parse_args([]) + assert (args.host, args.port) == ('0.0.0.0', 80) + + +def test_a_flag_beats_everything(): + args = _parser().parse_args(['--weight-type', 'FP32']) + assert args.weight_type == 'FP32' + + +def test_a_setting_is_read_from_the_variable_named_after_its_flag(monkeypatch: pytest.MonkeyPatch): + monkeypatch.setenv('WEIGHT_TYPE', 'FP32') + assert _parser().parse_args([]).weight_type == 'FP32' + + +def test_the_loop_own_host_is_never_mistaken_for_the_bind_address(monkeypatch: pytest.MonkeyPatch): + """HOST is the loop's address. Binding uvicorn to it would leave the node unreachable.""" + monkeypatch.setenv('HOST', 'preview.learning-loop.ai') + assert _parser().parse_args([]).host == '0.0.0.0' + + +def test_the_bind_address_has_a_name_of_its_own(monkeypatch: pytest.MonkeyPatch): + monkeypatch.setenv('NODE_HOST', '127.0.0.1') + monkeypatch.setenv('NODE_PORT', '8080') + args = _parser().parse_args([]) + assert (args.host, args.port) == ('127.0.0.1', 8080) + + +def test_a_node_that_used_a_prefix_still_reads_it(monkeypatch: pytest.MonkeyPatch): + monkeypatch.setenv('DFINE_DETECTOR_WEIGHT_TYPE', 'FP32') + parser = _parser(legacy_env_prefix='DFINE_DETECTOR_') + assert parser.parse_args([]).weight_type == 'FP32' + + +def test_the_prefixed_name_warns_which_one_to_use_instead(monkeypatch: pytest.MonkeyPatch, + caplog: pytest.LogCaptureFixture): + monkeypatch.setenv('DFINE_DETECTOR_WEIGHT_TYPE', 'FP32') + _parser(legacy_env_prefix='DFINE_DETECTOR_').parse_args([]) + assert 'DFINE_DETECTOR_WEIGHT_TYPE' in caplog.text + assert 'WEIGHT_TYPE' in caplog.text + + +def test_the_current_name_wins_over_the_prefixed_one(monkeypatch: pytest.MonkeyPatch): + monkeypatch.setenv('DFINE_DETECTOR_WEIGHT_TYPE', 'FP32') + monkeypatch.setenv('WEIGHT_TYPE', 'FP16') + assert _parser(legacy_env_prefix='DFINE_DETECTOR_').parse_args([]).weight_type == 'FP16' + + +def test_a_node_without_a_legacy_prefix_ignores_prefixed_names(monkeypatch: pytest.MonkeyPatch): + monkeypatch.setenv('DFINE_DETECTOR_WEIGHT_TYPE', 'FP32') + assert _parser().parse_args([]).weight_type == 'FP16' diff --git a/pyproject.toml b/pyproject.toml index 47409dff..d05073f5 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -13,6 +13,7 @@ dependencies = [ "python-socketio>=5.16.2,<6.0.0", "aiofiles>=0.7.0", "python-multipart>=0.0.31", + "configargparse>=1.7.1", "psutil>=5.9.0,<8.0.0", "numpy>=2.0,<3.0", "Pillow>=12.3.0,<13.0.0", diff --git a/uv.lock b/uv.lock index da0c0bf7..1b373e19 100644 --- a/uv.lock +++ b/uv.lock @@ -449,6 +449,15 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/d1/d6/3965ed04c63042e047cb6a3e6ed1a63a35087b6a609aa3a15ed8ac56c221/colorama-0.4.6-py2.py3-none-any.whl", hash = "sha256:4f1d9991f5acc0ca119f9d443620b77f9d6b33703e51011c16baf57afb285fc6", size = 25335, upload-time = "2022-10-25T02:36:20.889Z" }, ] +[[package]] +name = "configargparse" +version = "1.7.5" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/3f/0b/30328302903c55218ffc5199646d0e9d28348ff26c02ba77b2ffc58d294a/configargparse-1.7.5.tar.gz", hash = "sha256:e3f9a7bb6be34d66b2e3c4a2f58e3045f8dfae47b0dc039f87bcfaa0f193fb0f", size = 53548, upload-time = "2026-03-11T02:19:38.144Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/fe/19/3ba5e1b0bcc7b91aeab6c258afd70e4907d220fed3972febe38feb40db30/configargparse-1.7.5-py3-none-any.whl", hash = "sha256:1e63fdffedf94da9cd435fc13a1cd24777e76879dd2343912c1f871d4ac8c592", size = 27692, upload-time = "2026-03-11T02:19:36.442Z" }, +] + [[package]] name = "dacite" version = "1.9.2" @@ -781,6 +790,7 @@ source = { virtual = "." } dependencies = [ { name = "aiofiles" }, { name = "aiohttp" }, + { name = "configargparse" }, { name = "dacite" }, { name = "fastapi" }, { name = "httpx" }, @@ -813,6 +823,7 @@ requires-dist = [ { name = "aiofiles", specifier = ">=0.7.0" }, { name = "aiohttp", specifier = ">=3.14.3,<4.0.0" }, { name = "autopep8", marker = "extra == 'dev'", specifier = ">=2.0.2,<3.0.0" }, + { name = "configargparse", specifier = ">=1.7.1" }, { name = "dacite", specifier = ">=1.8.1,<2.0.0" }, { name = "debugpy", marker = "extra == 'dev'", specifier = ">=1.6.7.post1,<2.0.0" }, { name = "fastapi", specifier = ">=0.135,<1.0" }, From f3b3ca94accd6e84611feff76a2ab7bb022aae97 Mon Sep 17 00:00:00 2001 From: Kevin Heye Date: Thu, 20 Aug 2026 19:52:49 +0200 Subject: [PATCH 06/29] Share the framework-agnostic parts of a trainer Three things every trainer needs, none of which depend on the training framework, all of which existed only inside dfine_node or in worse form elsewhere: iterator_cpu_bound runs a training generator in a spawned process. A trainer that trains in-process needs this - the event loop must stay responsive and CUDA state must stay out of the node - and it is pure stdlib. find_batch_size probes for the largest power-of-two batch that fits, around a fits predicate the trainer supplies. The search is pure arithmetic; only the predicate needs a framework, so nothing here imports torch. Three repositories had three different implementations of this: a real probe, an estimate parsed out of torchinfo's summary text, and a doubling loop catching bare RuntimeError. is_out_of_memory exists because of that last one - catching RuntimeError treats every crash as a full card. macro_f1 scores the confusion matrix the loop stores, which was implemented three times over two different data shapes. _get_executor_error_from_log gets a default implementation, because yolov5_node and classification_node had byte-identical copies of it. The PEP 695 type parameters the generator used are rewritten as TypeVar and ParamSpec: this library supports Python 3.10 and dfine_node targets 3.12. Co-Authored-By: Claude Opus 5 (1M context) --- AGENTS.md | 9 ++ .../tests/unit/test_batch_size.py | 118 ++++++++++++++++++ learning_loop_node/tests/unit/test_metrics.py | 39 ++++++ .../tests/unit/test_subprocess.py | 54 ++++++++ learning_loop_node/trainer/batch_size.py | 86 +++++++++++++ learning_loop_node/trainer/metrics.py | 36 ++++++ learning_loop_node/trainer/subprocess.py | 93 ++++++++++++++ learning_loop_node/trainer/trainer_logic.py | 14 ++- 8 files changed, 447 insertions(+), 2 deletions(-) create mode 100644 learning_loop_node/tests/unit/test_batch_size.py create mode 100644 learning_loop_node/tests/unit/test_metrics.py create mode 100644 learning_loop_node/tests/unit/test_subprocess.py create mode 100644 learning_loop_node/trainer/batch_size.py create mode 100644 learning_loop_node/trainer/metrics.py create mode 100644 learning_loop_node/trainer/subprocess.py diff --git a/AGENTS.md b/AGENTS.md index c7877d5d..b6158d96 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -54,6 +54,15 @@ on 401/429) for the REST API, and a socket.io *client* for status updates and lo - **Annotator** — thin: it forwards the loop frontend's `handle_user_input` events into `AnnotatorLogic` and keeps a per-frontend history. +`trainer/subprocess.py`, `trainer/batch_size.py` and `trainer/metrics.py` hold the parts of a +trainer that are not framework-specific. `iterator_cpu_bound` runs a training generator in a +spawned process and yields its progress through a `maxsize=1` queue, so the event loop stays +responsive, CUDA state stays out of the node process, and the training can never run more than +one item ahead of the bookkeeping. `find_batch_size` probes for the largest power-of-two batch +that fits, around a `fits` predicate the trainer supplies — none of it imports torch, so a node +brings its own way of running a step. `macro_f1` scores the confusion matrix +`_get_new_best_training_state` returns. + `helpers/entrypoint.py` holds what every node's `main.py` repeats: `node_parser` builds a configargparse parser with `--host`/`--port`, `run_node` starts uvicorn. A setting is a flag *and* an environment variable from one declaration — `--conf-threshold` reads diff --git a/learning_loop_node/tests/unit/test_batch_size.py b/learning_loop_node/tests/unit/test_batch_size.py new file mode 100644 index 00000000..6b499a6a --- /dev/null +++ b/learning_loop_node/tests/unit/test_batch_size.py @@ -0,0 +1,118 @@ +from collections.abc import Callable + +import pytest + +from ...trainer.batch_size import ( + MIN_TRAIN_STEPS_PER_EPOCH, + batch_count, + dataset_limit, + find_batch_size, + is_out_of_memory, + no_gpu_batch_size, + smaller_pot, +) + + +def _recording(capacity: int) -> tuple[Callable[[int], bool], list[int]]: + """A `fits` predicate for a machine of `capacity`, plus the sizes it gets asked about.""" + calls: list[int] = [] + + def fits(batch_size: int) -> bool: + calls.append(batch_size) + return batch_size <= capacity + + return fits, calls + + +def _fits_up_to(capacity: int) -> Callable[[int], bool]: + fits, _ = _recording(capacity) + return fits + + +def test_the_search_doubles_up_to_the_limit(): + fits, calls = _recording(1024) + assert find_batch_size(fits, limit=16) == 16 + assert calls == [1, 2, 4, 8, 16] + + +def test_the_search_backs_off_to_the_last_size_that_fit(): + fits, calls = _recording(20) + assert find_batch_size(fits, limit=64) == 16 + assert calls == [1, 2, 4, 8, 16, 32] # probes 32, fails, keeps 16 + + +def test_a_limit_that_is_not_a_power_of_two_is_rounded_down(): + assert find_batch_size(_fits_up_to(1024), limit=48) == 32 + assert find_batch_size(_fits_up_to(1024), limit=1) == 1 + + +def test_a_machine_that_cannot_take_one_sample_is_an_error(): + fits, calls = _recording(0) + with pytest.raises(RuntimeError, match='batch size 1 does not fit'): + find_batch_size(fits, limit=64) + assert calls == [1], 'must give up instead of probing larger sizes' + + +@pytest.mark.parametrize('capacity', range(1, 130)) +def test_the_result_is_always_the_largest_power_of_two_that_fits(capacity: int): + assert find_batch_size(_fits_up_to(capacity), limit=512) == smaller_pot(capacity) + + +def test_equal_hardware_yields_an_equal_recipe(): + """Only powers of two, so two machines of similar size train identically.""" + assert find_batch_size(_fits_up_to(37), limit=512) == find_batch_size(_fits_up_to(39), limit=512) + + +def test_smaller_pot(): + assert [smaller_pot(n) for n in (1, 2, 3, 4, 7, 8, 15, 1293)] == [1, 2, 2, 4, 4, 8, 8, 1024] + with pytest.raises(ValueError, match='n must be >= 1'): + smaller_pot(0) + + +def test_batch_count_covers_the_whole_set_without_overshooting_by_a_batch(): + for sample_count in (1, 7, 8, 900, 5000): + for batch_size in (1, 2, 8, 64, 512): + covered = batch_count(sample_count, batch_size) * batch_size + assert covered >= sample_count + assert covered - sample_count < batch_size + + +def test_the_dataset_limit_keeps_enough_steps_per_epoch(): + for sample_count in (8, 20, 47, 100, 1000, 118_000): + limit = smaller_pot(dataset_limit(sample_count)) + assert sample_count // limit >= MIN_TRAIN_STEPS_PER_EPOCH + + +def test_the_dataset_limit_stays_usable_for_a_tiny_set(): + assert dataset_limit(7) == 1 + with pytest.raises(ValueError, match='sample_count must be >= 1'): + dataset_limit(0) + + +def test_memory_still_decides_below_the_dataset_limit(): + """The bound is a ceiling only: a card that fits just 2 keeps training at 2.""" + assert find_batch_size(_fits_up_to(2), limit=dataset_limit(20)) == 2 + + +def test_without_a_gpu_the_fallback_respects_the_limit(): + assert no_gpu_batch_size(1024, 'probe') == 8 + assert no_gpu_batch_size(2, 'probe') == 2 + + +@pytest.mark.parametrize('message', [ + 'CUDA out of memory. Tried to allocate 20.00 MiB', + 'cuDNN error: CUDNN_STATUS_ALLOC_FAILED', + 'CUDA error: unknown error', +]) +def test_allocation_failures_are_recognised_however_they_surface(message: str): + assert is_out_of_memory(RuntimeError(message)) + + +def test_a_real_bug_is_not_mistaken_for_a_full_card(): + """A trainer catching bare RuntimeError treats every crash as 'too big'; this does not.""" + assert not is_out_of_memory(RuntimeError('shape mismatch in forward pass')) + assert not is_out_of_memory(ValueError('bad config')) + + +def test_the_host_running_out_of_memory_counts_too(): + assert is_out_of_memory(MemoryError()) diff --git a/learning_loop_node/tests/unit/test_metrics.py b/learning_loop_node/tests/unit/test_metrics.py new file mode 100644 index 00000000..7c8537dc --- /dev/null +++ b/learning_loop_node/tests/unit/test_metrics.py @@ -0,0 +1,39 @@ +import pytest + +from ...trainer.metrics import category_f1, confusion_matrix_from_counts, macro_f1 + + +def test_the_score_averages_the_categories_instead_of_pooling_them(): + """Counts from a real epoch: face 248/19/59 and hand 39/54/73 - 74% pooled, 62% averaged.""" + assert macro_f1({'face': {'tp': 248, 'fp': 19, 'fn': 59}, + 'hand': {'tp': 39, 'fp': 54, 'fn': 73}}) == pytest.approx(0.622, abs=0.001) + assert macro_f1({'both': {'tp': 287, 'fp': 73, 'fn': 132}}) == pytest.approx(0.737, abs=0.001) + + +def test_a_rare_category_carries_the_same_weight(): + """A category the model never finds halves the score, however few instances it has.""" + assert macro_f1({'frequent': {'tp': 1000, 'fp': 0, 'fn': 0}, + 'rare': {'tp': 0, 'fp': 0, 'fn': 3}}) == pytest.approx(0.5) + + +def test_a_category_without_any_counts_scores_zero_instead_of_raising(): + assert macro_f1({'unseen': {'tp': 0, 'fp': 0, 'fn': 0}}) == 0.0 + assert category_f1({'tp': 0, 'fp': 0, 'fn': 0}) == 0.0 + + +def test_an_empty_confusion_matrix_scores_zero(): + assert macro_f1({}) == 0.0 + + +def test_a_perfect_category_scores_one(): + assert category_f1({'tp': 10, 'fp': 0, 'fn': 0}) == pytest.approx(1.0) + + +def test_precision_and_recall_are_balanced(): + assert category_f1({'tp': 5, 'fp': 5, 'fn': 0}) == category_f1({'tp': 5, 'fp': 0, 'fn': 5}) + + +def test_counts_can_be_assembled_from_separate_mappings(): + matrix = confusion_matrix_from_counts({'a': 3}, {'a': 1, 'b': 2}, {'b': 4}) + assert matrix == {'a': {'tp': 3, 'fp': 1, 'fn': 0}, + 'b': {'tp': 0, 'fp': 2, 'fn': 4}} diff --git a/learning_loop_node/tests/unit/test_subprocess.py b/learning_loop_node/tests/unit/test_subprocess.py new file mode 100644 index 00000000..f2517ced --- /dev/null +++ b/learning_loop_node/tests/unit/test_subprocess.py @@ -0,0 +1,54 @@ +import pytest + +from ...trainer.subprocess import iterator_cpu_bound + + +def counting(limit: int): + yield from range(limit) + + +def raising_after(limit: int): + yield from range(limit) + raise RuntimeError('the training crashed') + + +def empty(): + return iter(()) + + +async def _collect(iterator) -> list: + return [item async for item in iterator] + + +async def test_everything_the_generator_yields_arrives_in_order(): + async with iterator_cpu_bound(counting, 5) as iterator: + assert await _collect(iterator) == [0, 1, 2, 3, 4] + + +async def test_a_generator_that_yields_nothing_simply_finishes(): + async with iterator_cpu_bound(empty) as iterator: + assert await _collect(iterator) == [] + + +async def test_a_failure_in_the_process_is_raised_in_the_caller(): + """Otherwise a crashed training would look like a training that finished.""" + with pytest.raises(RuntimeError, match='the training crashed'): + async with iterator_cpu_bound(raising_after, 2) as iterator: + await _collect(iterator) + + +async def test_what_the_generator_produced_before_failing_still_arrives(): + received = [] + with pytest.raises(RuntimeError): + async with iterator_cpu_bound(raising_after, 3) as iterator: + async for item in iterator: + received.append(item) + assert received == [0, 1, 2] + + +async def test_leaving_early_does_not_leave_the_process_running(): + async with iterator_cpu_bound(counting, 1000) as iterator: + async for item in iterator: + if item == 2: + break + # the context manager killed and joined the process; reaching here without hanging is the test diff --git a/learning_loop_node/trainer/batch_size.py b/learning_loop_node/trainer/batch_size.py new file mode 100644 index 00000000..707bb8c6 --- /dev/null +++ b/learning_loop_node/trainer/batch_size.py @@ -0,0 +1,86 @@ +"""Choosing a batch size by probing, rather than configuring one. + +The largest batch that fits depends on the model, the image resolution and the card, so a +configured value is either wrong on some machines or too small on all of them. Probing answers +it directly: run a representative step at doubling sizes and keep the last one that survived. + +What "a representative step" means is the trainer's business — it supplies a ``fits`` predicate +that runs a real step and reports whether it ran out of memory. Everything here is that +predicate's scaffolding, and none of it imports a deep-learning framework. + +Only powers of two are visited, so equal hardware yields an equal recipe. That matters when a +trainer scales its learning rate by the batch size: a probe that returned 37 on one machine and +39 on another would make the two trainings quietly different. + +Adapted from PyTorch Lightning's ``BatchSizeFinder`` (power-scaling mode). +Copyright The Lightning AI team. Licensed under the Apache License, Version 2.0. +https://github.com/Lightning-AI/pytorch-lightning +""" + +import logging +from collections.abc import Callable + +logger = logging.getLogger(__name__) + +NO_GPU_BATCH_SIZE = 8 +"""What to fall back on with no GPU to probe: enough to make progress, small enough to fit.""" + +MIN_TRAIN_STEPS_PER_EPOCH = 8 +"""Below this an epoch is one or two optimizer steps, which trains poorly whatever the card.""" + + +def find_batch_size(fits: Callable[[int], bool], *, limit: int) -> int: + """Return the largest power-of-two batch size that fits, never exceeding ``limit``. + + :param fits: Runs a representative probe; ``False`` on out-of-memory. + :raises RuntimeError: If not even a batch size of 1 fits. + """ + limit = smaller_pot(limit) + if not fits(1): + raise RuntimeError('batch size 1 does not fit in memory') + + size = 1 + while size < limit and fits(size * 2): + size *= 2 + + return size + + +def dataset_limit(sample_count: int) -> int: + """The batch size ceiling that still leaves ``MIN_TRAIN_STEPS_PER_EPOCH`` steps per epoch.""" + if sample_count < 1: + raise ValueError(f'sample_count must be >= 1, got {sample_count}') + return max(1, sample_count // MIN_TRAIN_STEPS_PER_EPOCH) + + +def no_gpu_batch_size(limit: int, probe: str) -> int: + """The batch size to fall back on when there is no GPU to probe.""" + batch_size = min(smaller_pot(limit), NO_GPU_BATCH_SIZE) + logger.warning('%s: CUDA is unavailable; using batch size %d without probing', probe, batch_size) + return batch_size + + +def smaller_pot(n: int) -> int: + """The largest power of two that is <= ``n``.""" + if n < 1: + raise ValueError(f'n must be >= 1, got {n}') + return 1 << (n.bit_length() - 1) + + +def batch_count(sample_count: int, batch_size: int) -> int: + """How many batches a loader yields, the last one possibly short.""" + return -(-sample_count // batch_size) + + +def is_out_of_memory(exception: BaseException) -> bool: + """Whether the exception signals exhausted memory, on the GPU or the host. + + Allocation failures do not all surface as a framework's dedicated error type: cuDNN and + cuBLAS workspaces raise a plain ``RuntimeError``. Catching those by message is what keeps a + probe from mistaking a real bug for a full card — a trainer that catches bare + ``RuntimeError`` around its probe silently treats every crash as "too big". + """ + if isinstance(exception, MemoryError): + return True + message = str(exception).lower() + return any(text in message for text in ('out of memory', 'alloc_failed', 'cuda error: unknown error')) diff --git a/learning_loop_node/trainer/metrics.py b/learning_loop_node/trainer/metrics.py new file mode 100644 index 00000000..aee553b8 --- /dev/null +++ b/learning_loop_node/trainer/metrics.py @@ -0,0 +1,36 @@ +"""Scoring a training from the confusion matrix the loop stores. + +The loop keeps one ``{'tp': .., 'fp': .., 'fn': ..}`` per category, which is what +``TrainerLogicGeneric._get_new_best_training_state`` returns. Deciding whether an epoch beat +the previous best means reducing that to one number, and every trainer needs the same one. +""" + +import statistics +from collections.abc import Mapping + + +def macro_f1(confusion_matrix: Mapping[str, Mapping[str, int]]) -> float: + """The unweighted mean of the per-category F1 scores, as the loop's UI shows it by default. + + Unweighted is the point: a rare category the model never finds drags the score down as much + as a frequent one, where pooling the counts first would hide it. + """ + scores = [category_f1(counts) for counts in confusion_matrix.values()] + return statistics.mean(scores) if scores else 0.0 + + +def category_f1(counts: Mapping[str, int]) -> float: + """The F1 score of one category, 0.0 where it is undefined rather than a division error.""" + tp, fp, fn = counts['tp'], counts['fp'], counts['fn'] + precision = tp / (tp + fp) if tp + fp else 0.0 + recall = tp / (tp + fn) if tp + fn else 0.0 + return 2 * precision * recall / (precision + recall) if precision + recall else 0.0 + + +def confusion_matrix_from_counts(true_positives: Mapping[str, int], false_positives: Mapping[str, int], + false_negatives: Mapping[str, int]) -> dict[str, dict[str, int]]: + """Assemble the loop's confusion matrix from three per-category count mappings.""" + return {category: {'tp': true_positives.get(category, 0), + 'fp': false_positives.get(category, 0), + 'fn': false_negatives.get(category, 0)} + for category in {*true_positives, *false_positives, *false_negatives}} diff --git a/learning_loop_node/trainer/subprocess.py b/learning_loop_node/trainer/subprocess.py new file mode 100644 index 00000000..dbfa181b --- /dev/null +++ b/learning_loop_node/trainer/subprocess.py @@ -0,0 +1,93 @@ +"""Run a blocking, CPU-bound generator in its own process without blocking the event loop. + +A trainer that trains in-process has a problem: the training must not stall the node, and CUDA +state must stay out of the node process so a crashed training cannot take the node with it. +:func:`iterator_cpu_bound` runs the generator in a spawned process and yields what it produces +through a ``maxsize=1`` queue, so the producer can never run more than one item ahead of the +bookkeeping that consumes it — which is what lets a trainer alternate between two model files +and know the one it is copying is not being rewritten. + +Exceptions raised inside the process are re-raised in the caller, and the process is killed if +the caller leaves the context early. +""" +from __future__ import annotations + +import asyncio +import multiprocessing +import queue +import traceback +from collections.abc import AsyncGenerator, Callable, Iterator +from contextlib import asynccontextmanager +from multiprocessing.queues import Queue as MPQueue +from typing import Any, ParamSpec, TypeVar + +T = TypeVar('T') +P = ParamSpec('P') + + +class IteratorDone: + pass + + +def _iterator_wrapper( + it: Callable[..., Iterator[T]], + state_queue: MPQueue[T | Exception | IteratorDone], + args: tuple[Any, ...], + kwargs: dict[str, Any], +) -> None: + try: + for data in it(*args, **kwargs): + state_queue.put(data) + except Exception as e: + print(traceback.format_exc()) + state_queue.put(e) + + state_queue.put(IteratorDone()) + + +async def _iterator_cpu_bound_inner( + it: Callable[P, Iterator[T]], + *args: P.args, + **kwargs: P.kwargs, +) -> AsyncGenerator[T, None]: + state_queue: MPQueue[T | Exception | IteratorDone] = multiprocessing.Queue(maxsize=1) + process = multiprocessing.Process( + target=_iterator_wrapper, + args=(it, state_queue, args, kwargs), + name='iterator_cpu_bound', + ) + + process.start() + + try: + while True: + try: + item = await asyncio.to_thread(state_queue.get, True, 0.5) + except queue.Empty: + if not process.is_alive(): + break + continue + match item: + case IteratorDone(): + break + case Exception() as e: + raise e + case _ as other: + yield other + finally: + if process.is_alive(): + process.kill() + process.join() + + +@asynccontextmanager +async def iterator_cpu_bound( + it: Callable[P, Iterator[T]], + *args: P.args, + **kwargs: P.kwargs, +) -> AsyncGenerator[AsyncGenerator[T, None], None]: + iterator = _iterator_cpu_bound_inner(it, *args, **kwargs) + try: + yield iterator + finally: + await asyncio.shield(iterator.aclose()) diff --git a/learning_loop_node/trainer/trainer_logic.py b/learning_loop_node/trainer/trainer_logic.py index c5a9bade..9827b41b 100644 --- a/learning_loop_node/trainer/trainer_logic.py +++ b/learning_loop_node/trainer/trainer_logic.py @@ -178,9 +178,19 @@ async def _resume(self) -> None: '''Is called when self.can_resume() returns True. One may resume the training on a previously trained model stored by self.on_model_published(basic_model).''' - @abstractmethod def _get_executor_error_from_log(self) -> Optional[str]: - '''Should be used to provide error informations to the Learning Loop by extracting data from self.executor.get_log().''' + '''Reports what went wrong to the Learning Loop by reading self.executor's log. + + The default recognises the CUDA failures every trainer hits. Override to add messages a + particular training framework produces, and call super() to keep these.''' + if self._executor is None: + return None + for line in self._executor.get_log_by_lines(tail=50): + if 'CUDA out of memory' in line: + return 'graphics card is out of memory' + if 'CUDA error: invalid device ordinal' in line: + return 'graphics card not found' + return None @abstractmethod async def _detect(self, model_information: ModelInformation, images: List[str], model_folder: str) -> List[Detections]: From 13c7793bb4f0107e66b95c45b2b9e4435172d381 Mon Sep 17 00:00:00 2001 From: Kevin Heye Date: Tue, 25 Aug 2026 11:30:50 +0200 Subject: [PATCH 07/29] Resolve a disagreeing environment variable instead of giving up read_from_env accepts an ordered list of names - ['LOOP_HOST', 'HOST'] and so on - and returned None when two of them were set to different values. The caller then fell back to its default, and host()'s default is 'learning-loop.ai': a .env carrying HOST and LOOP_HOST with different values pointed the node at production, with only a warning in the log to say so. Since values is built in the order of possible_names and then filtered, the first surviving value is already the one the list prefers - so the fix is simply to warn and fall through to the existing return instead of bailing out. ignore_errors=False still raises, as before. This is the right place for the fix. The node repositories' docker.sh scripts had been working around it by forwarding exactly one spelling, which meant every sub-project's list had to agree on which one; with the container receiving the whole .env, only the library can decide. Adds tests/unit/test_environment_reader.py: the prefixed name wins, either name alone is read, a disagreement resolves to the preferred name rather than the default, and ignore_errors=False still raises. Co-Authored-By: Claude Opus 5 --- .../helpers/environment_reader.py | 13 +++-- .../tests/unit/test_environment_reader.py | 51 +++++++++++++++++++ 2 files changed, 59 insertions(+), 5 deletions(-) create mode 100644 learning_loop_node/tests/unit/test_environment_reader.py diff --git a/learning_loop_node/helpers/environment_reader.py b/learning_loop_node/helpers/environment_reader.py index 4833e1b2..3cb8cca3 100644 --- a/learning_loop_node/helpers/environment_reader.py +++ b/learning_loop_node/helpers/environment_reader.py @@ -15,12 +15,15 @@ def read_from_env(possible_names: List[str], ignore_errors: bool = True) -> Opti return None raise ValueError(f'no environment variable set for {possible_names}') - # Possible error: multiple values are not None and not equal + # Possible error: multiple values are not None and not equal. + # NOTE returning None here would be worse than picking one: the caller falls back to its + # default, and host()'s default is the production loop. `possible_names` is ordered by + # preference, so a disagreement resolves to the first name that is set. if len(values) > 1 and len(set(values)) > 1: - if ignore_errors: - logging.warning('different environment variables set for %s: %s', possible_names, values) - return None - raise ValueError(f'different environment variables set for {possible_names}: {values}') + if not ignore_errors: + raise ValueError(f'different environment variables set for {possible_names}: {values}') + logging.warning('different environment variables set for %s: %s - using %s', + possible_names, values, values[0]) return values[0] diff --git a/learning_loop_node/tests/unit/test_environment_reader.py b/learning_loop_node/tests/unit/test_environment_reader.py new file mode 100644 index 00000000..a66a844e --- /dev/null +++ b/learning_loop_node/tests/unit/test_environment_reader.py @@ -0,0 +1,51 @@ +"""Tests for `helpers/environment_reader`, which resolves every loop setting a node reads.""" + +import pytest + +from ...helpers import environment_reader + +NAMES = ('LOOP_HOST', 'HOST', 'LOOP_ORGANIZATION', 'ORGANIZATION') + + +@pytest.fixture(autouse=True) +def clean_env(monkeypatch: pytest.MonkeyPatch): + for name in NAMES: + monkeypatch.delenv(name, raising=False) + + +def test_the_prefixed_name_is_preferred(monkeypatch: pytest.MonkeyPatch): + monkeypatch.setenv('LOOP_HOST', 'preview.learning-loop.ai') + monkeypatch.setenv('HOST', 'preview.learning-loop.ai') + assert environment_reader.host() == 'preview.learning-loop.ai' + + +def test_either_name_alone_is_read(monkeypatch: pytest.MonkeyPatch): + monkeypatch.setenv('HOST', 'preview.learning-loop.ai') + assert environment_reader.host() == 'preview.learning-loop.ai' + monkeypatch.delenv('HOST') + monkeypatch.setenv('LOOP_HOST', 'other.learning-loop.ai') + assert environment_reader.host() == 'other.learning-loop.ai' + + +def test_a_disagreement_resolves_to_the_preferred_name(monkeypatch: pytest.MonkeyPatch): + """Returning nothing here would let host() fall back to its default, which is production.""" + monkeypatch.setenv('LOOP_HOST', 'preview.learning-loop.ai') + monkeypatch.setenv('HOST', 'learning-loop.ai') + assert environment_reader.host(default='learning-loop.ai') == 'preview.learning-loop.ai' + + +def test_a_disagreement_falls_back_to_the_second_name_when_the_first_is_unset( + monkeypatch: pytest.MonkeyPatch): + monkeypatch.setenv('ORGANIZATION', 'zauberzeug') + assert environment_reader.organization() == 'zauberzeug' + + +def test_a_disagreement_still_raises_when_errors_are_not_ignored(monkeypatch: pytest.MonkeyPatch): + monkeypatch.setenv('LOOP_HOST', 'a') + monkeypatch.setenv('HOST', 'b') + with pytest.raises(ValueError, match='different environment variables'): + environment_reader.read_from_env(['LOOP_HOST', 'HOST'], ignore_errors=False) + + +def test_nothing_set_yields_the_default(): + assert environment_reader.host(default='fallback') == 'fallback' From 8e402d6d6a91faea5957509c5ac8b6f4b46ab63b Mon Sep 17 00:00:00 2001 From: Kevin Heye Date: Tue, 25 Aug 2026 13:24:27 +0200 Subject: [PATCH 08/29] Document what a disagreeing alias resolves to The environment variable table already lists the alias for each loop setting but said nothing about both being set. That is the case the previous commit changed: the prefixed name wins, and the variable is never treated as unset, which would have let LOOP_HOST fall back to learning-loop.ai. Co-Authored-By: Claude Opus 5 --- README.md | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/README.md b/README.md index 6372a5b3..0297b05f 100644 --- a/README.md +++ b/README.md @@ -13,6 +13,8 @@ This Python library helps to write Nodes that interact with the Zauberzeug Learn To start a node you have to implement the logic by inheriting from the corresponding base logic class. We provide samples in the 'mock' folders and recommend to follow that scheme. A complete trainer and detector example can be found [here](https://github.com/zauberzeug/yolov5_node). +**[docs/writing-a-node.md](docs/writing-a-node.md) walks through building a node from an empty repository**: which base class to implement per node type, what the library does for you, the trainer state machine, the trainer-to-detector model contract, and the repository and docker conventions the existing nodes share. The rest of this README documents how to *operate* a node. + #### Environment variables You can configure connection to our Learning Loop by specifying the following environment variables before starting: @@ -36,6 +38,10 @@ You can configure connection to our Learning Loop by specifying the following en Note that organization and project IDs are always lower case and may differ from the names in the Learning Loop which can have uppercase letters. +Where a name has an alias, either spelling works. If both are set to **different** values the +prefixed name wins and a warning names the value used — the variable is never treated as unset, +which would otherwise let `LOOP_HOST` fall back to its default of `learning-loop.ai`. + #### Testing We use github actions for CI. Tests can also be executed locally by running From d358da5b6463501c09df29d5328ec6a4fec6fffc Mon Sep 17 00:00:00 2001 From: Kevin Heye Date: Wed, 26 Aug 2026 11:40:27 +0200 Subject: [PATCH 09/29] Move the category lookups out of postprocess category_by_index and category_by_name resolve a model's metadata; they do not post-process anything. It showed at the one call site outside this library: classification_node's detector, which produces no boxes at all, imported category_by_name from a module called postprocess. They move to detector/categories.py, which leaves postprocess to the job its name describes - suppression, detection building and conversion to the loop's dataclasses. Both functions are unchanged, and their tests move with them. Worth doing now rather than later: these are public API the moment the library is released, and renaming a module afterwards is a breaking change. Co-Authored-By: Claude Opus 5 (1M context) --- learning_loop_node/detector/categories.py | 37 +++++++++++++++++++ learning_loop_node/detector/postprocess.py | 30 +-------------- .../tests/unit/test_categories.py | 33 +++++++++++++++++ .../tests/unit/test_postprocess.py | 24 ------------ 4 files changed, 71 insertions(+), 53 deletions(-) create mode 100644 learning_loop_node/detector/categories.py create mode 100644 learning_loop_node/tests/unit/test_categories.py diff --git a/learning_loop_node/detector/categories.py b/learning_loop_node/detector/categories.py new file mode 100644 index 00000000..9f45541f --- /dev/null +++ b/learning_loop_node/detector/categories.py @@ -0,0 +1,37 @@ +"""Resolving a model's categories from what it emits. + +A model reports class indices, or names, and only ``ModelInformation.categories`` gives those +meaning. Both lookups are the point at which a model and its metadata are checked against each +other, so both raise rather than skipping quietly: a mismatch here makes every prediction on the +image suspect, and it is far cheaper to diagnose at the lookup than three layers later. +""" + +from ..data_classes import Category, ModelInformation + + +def category_by_index(model_information: ModelInformation, index: int) -> Category: + """Resolve the category a model's class index refers to. + + Models emit class indices, and the order of ``model_information.categories`` is what + gives them meaning — so an out-of-range index is a model/metadata mismatch, not a + detection to skip quietly. + + :raises ValueError: If the index is outside the model's category list. + """ + categories = model_information.categories + if not 0 <= index < len(categories): + raise ValueError( + f'category index {index} is out of range for a model with {len(categories)} categories') + return categories[index] + + +def category_by_name(model_information: ModelInformation, name: str) -> Category: + """Resolve a category by name, for models whose outputs are named rather than indexed. + + :raises ValueError: If no category of that name exists. + """ + for category in model_information.categories: + if category.name == name: + return category + known = ', '.join(category.name for category in model_information.categories) + raise ValueError(f'unknown category name {name!r}; the model knows: {known}') diff --git a/learning_loop_node/detector/postprocess.py b/learning_loop_node/detector/postprocess.py index ae47fbcc..8478af81 100644 --- a/learning_loop_node/detector/postprocess.py +++ b/learning_loop_node/detector/postprocess.py @@ -19,13 +19,13 @@ from ..data_classes import ( BoxDetection, - Category, Detections, ImageMetadata, ModelInformation, PointDetection, ) from ..enums import CategoryType +from .categories import category_by_index from .geometry import clip_box, clip_point MIN_BOX_SIZE: int = 2 @@ -36,34 +36,6 @@ into :attr:`ModelInformation.categories`.""" -def category_by_index(model_information: ModelInformation, index: int) -> Category: - """Resolve the category a model's class index refers to. - - Models emit class indices, and the order of ``model_information.categories`` is what - gives them meaning — so an out-of-range index is a model/metadata mismatch, not a - detection to skip quietly. - - :raises ValueError: If the index is outside the model's category list. - """ - categories = model_information.categories - if not 0 <= index < len(categories): - raise ValueError( - f'category index {index} is out of range for a model with {len(categories)} categories') - return categories[index] - - -def category_by_name(model_information: ModelInformation, name: str) -> Category: - """Resolve a category by name, for models whose outputs are named rather than indexed. - - :raises ValueError: If no category of that name exists. - """ - for category in model_information.categories: - if category.name == name: - return category - known = ', '.join(category.name for category in model_information.categories) - raise ValueError(f'unknown category name {name!r}; the model knows: {known}') - - def bbox_iou( box1: np.ndarray, box2: np.ndarray, diff --git a/learning_loop_node/tests/unit/test_categories.py b/learning_loop_node/tests/unit/test_categories.py new file mode 100644 index 00000000..4eacfebb --- /dev/null +++ b/learning_loop_node/tests/unit/test_categories.py @@ -0,0 +1,33 @@ +import pytest + +from ...data_classes import Category, ModelInformation +from ...detector.categories import category_by_index, category_by_name +from ...enums import CategoryType + +BOX = Category(id='uuid-box', name='car', type=CategoryType.Box) +POINT = Category(id='uuid-point', name='weed', type=CategoryType.Point) + + +def model_information(*categories: Category) -> ModelInformation: + return ModelInformation(id='model-uuid', host='localhost', organization='zauberzeug', + project='pytest', version='1.2', categories=list(categories or (BOX, POINT))) + + +def test_category_is_resolved_by_index(): + assert category_by_index(model_information(), 1) is POINT + + +@pytest.mark.parametrize('index', [-1, 2, 99]) +def test_an_index_outside_the_model_categories_is_an_error(index: int): + # a mismatch between model and metadata must not be silently skipped + with pytest.raises(ValueError, match='out of range'): + category_by_index(model_information(), index) + + +def test_category_is_resolved_by_name(): + assert category_by_name(model_information(), 'weed') is POINT + + +def test_an_unknown_category_name_lists_the_known_ones(): + with pytest.raises(ValueError, match='car, weed'): + category_by_name(model_information(), 'tractor') diff --git a/learning_loop_node/tests/unit/test_postprocess.py b/learning_loop_node/tests/unit/test_postprocess.py index 14dbcd6a..c70562a0 100644 --- a/learning_loop_node/tests/unit/test_postprocess.py +++ b/learning_loop_node/tests/unit/test_postprocess.py @@ -5,8 +5,6 @@ from ...detector.postprocess import ( Detection, bbox_iou, - category_by_index, - category_by_name, detections_from_xyxy, non_max_suppression, post_process, @@ -24,28 +22,6 @@ def model_information(*categories: Category) -> ModelInformation: project='pytest', version='1.2', categories=list(categories or (BOX, POINT))) -# ---------------------------------------------------------------- category resolution - -def test_category_is_resolved_by_index(): - assert category_by_index(model_information(), 1) is POINT - - -@pytest.mark.parametrize('index', [-1, 2, 99]) -def test_an_index_outside_the_model_categories_is_an_error(index: int): - # a mismatch between model and metadata must not be silently skipped - with pytest.raises(ValueError, match='out of range'): - category_by_index(model_information(), index) - - -def test_category_is_resolved_by_name(): - assert category_by_name(model_information(), 'weed') is POINT - - -def test_an_unknown_category_name_lists_the_known_ones(): - with pytest.raises(ValueError, match='car, weed'): - category_by_name(model_information(), 'tractor') - - # ---------------------------------------------------------------- iou and suppression def test_identical_boxes_have_an_iou_of_one(): From e19e448203601e307112686091c389b6dffe41be Mon Sep 17 00:00:00 2001 From: Kevin Heye Date: Wed, 26 Aug 2026 14:58:50 +0200 Subject: [PATCH 10/29] Use a generic prefix in the legacy_env_prefix example This repository is public and the example named a node that is not. The prefix a private node once required is not what the parameter is about, and the tests need any prefix at all, not that one. 203 unit tests pass, ruff clean. Co-Authored-By: Claude Opus 5 --- learning_loop_node/helpers/entrypoint.py | 2 +- .../tests/unit/test_entrypoint.py | 18 +++++++++--------- 2 files changed, 10 insertions(+), 10 deletions(-) diff --git a/learning_loop_node/helpers/entrypoint.py b/learning_loop_node/helpers/entrypoint.py index 830c89ec..1b360c76 100644 --- a/learning_loop_node/helpers/entrypoint.py +++ b/learning_loop_node/helpers/entrypoint.py @@ -25,7 +25,7 @@ def node_parser(*, description: str, legacy_env_prefix: str = '') -> configargpa """Build the parser for a node, pre-loaded with the settings every node has. :param legacy_env_prefix: A prefix an earlier version of this node required, e.g. - ``'DFINE_DETECTOR_'``. Prefixed names are still honoured, with a warning, so a + ``'MY_DETECTOR_'``. Prefixed names are still honoured, with a warning, so a deployment keeps working until it is updated. Leave empty for a node that has always read unprefixed names. """ diff --git a/learning_loop_node/tests/unit/test_entrypoint.py b/learning_loop_node/tests/unit/test_entrypoint.py index a478563b..d20a0911 100644 --- a/learning_loop_node/tests/unit/test_entrypoint.py +++ b/learning_loop_node/tests/unit/test_entrypoint.py @@ -2,7 +2,7 @@ from ...helpers.entrypoint import node_parser -MANAGED = ('WEIGHT_TYPE', 'DFINE_DETECTOR_WEIGHT_TYPE', 'HOST', 'NODE_HOST', 'NODE_PORT', 'PORT') +MANAGED = ('WEIGHT_TYPE', 'MY_DETECTOR_WEIGHT_TYPE', 'HOST', 'NODE_HOST', 'NODE_PORT', 'PORT') @pytest.fixture(autouse=True) @@ -47,25 +47,25 @@ def test_the_bind_address_has_a_name_of_its_own(monkeypatch: pytest.MonkeyPatch) def test_a_node_that_used_a_prefix_still_reads_it(monkeypatch: pytest.MonkeyPatch): - monkeypatch.setenv('DFINE_DETECTOR_WEIGHT_TYPE', 'FP32') - parser = _parser(legacy_env_prefix='DFINE_DETECTOR_') + monkeypatch.setenv('MY_DETECTOR_WEIGHT_TYPE', 'FP32') + parser = _parser(legacy_env_prefix='MY_DETECTOR_') assert parser.parse_args([]).weight_type == 'FP32' def test_the_prefixed_name_warns_which_one_to_use_instead(monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture): - monkeypatch.setenv('DFINE_DETECTOR_WEIGHT_TYPE', 'FP32') - _parser(legacy_env_prefix='DFINE_DETECTOR_').parse_args([]) - assert 'DFINE_DETECTOR_WEIGHT_TYPE' in caplog.text + monkeypatch.setenv('MY_DETECTOR_WEIGHT_TYPE', 'FP32') + _parser(legacy_env_prefix='MY_DETECTOR_').parse_args([]) + assert 'MY_DETECTOR_WEIGHT_TYPE' in caplog.text assert 'WEIGHT_TYPE' in caplog.text def test_the_current_name_wins_over_the_prefixed_one(monkeypatch: pytest.MonkeyPatch): - monkeypatch.setenv('DFINE_DETECTOR_WEIGHT_TYPE', 'FP32') + monkeypatch.setenv('MY_DETECTOR_WEIGHT_TYPE', 'FP32') monkeypatch.setenv('WEIGHT_TYPE', 'FP16') - assert _parser(legacy_env_prefix='DFINE_DETECTOR_').parse_args([]).weight_type == 'FP16' + assert _parser(legacy_env_prefix='MY_DETECTOR_').parse_args([]).weight_type == 'FP16' def test_a_node_without_a_legacy_prefix_ignores_prefixed_names(monkeypatch: pytest.MonkeyPatch): - monkeypatch.setenv('DFINE_DETECTOR_WEIGHT_TYPE', 'FP32') + monkeypatch.setenv('MY_DETECTOR_WEIGHT_TYPE', 'FP32') assert _parser().parse_args([]).weight_type == 'FP16' From b73af70acddfaf34247fc3647e1fe6008130e24a Mon Sep 17 00:00:00 2001 From: Kevin Heye Date: Wed, 26 Aug 2026 15:21:28 +0200 Subject: [PATCH 11/29] Put the public functions first in the two moved modules CONTRIBUTING rule 6 wants the main functions at the top and every helper below its callers, so a reader meets the caller before the callee. Both modules came from node repositories the other way round: `subprocess.py` had the public `iterator_cpu_bound` last, below both helpers it calls, and `postprocess.py` had `_append_detections` above `to_image_metadata` and `to_detections`, which are its only two callers, while the NMS chain read callee-first throughout. Pure moves, no line changed - `git show --color-moved` shows it as one. The reorder was applied by script and asserted to leave the set of lines identical. 203 unit tests pass. Co-Authored-By: Claude Opus 5 --- learning_loop_node/detector/postprocess.py | 138 ++++++++++----------- learning_loop_node/trainer/subprocess.py | 54 ++++---- 2 files changed, 96 insertions(+), 96 deletions(-) diff --git a/learning_loop_node/detector/postprocess.py b/learning_loop_node/detector/postprocess.py index 8478af81..0fb509fa 100644 --- a/learning_loop_node/detector/postprocess.py +++ b/learning_loop_node/detector/postprocess.py @@ -36,24 +36,37 @@ into :attr:`ModelInformation.categories`.""" -def bbox_iou( - box1: np.ndarray, - box2: np.ndarray, -) -> np.ndarray: - """Compute IoU between box1 (1x4) and box2 (Nx4), both in x1y1x2y2 format.""" - b1_x1, b1_y1, b1_x2, b1_y2 = box1[:, 0], box1[:, 1], box1[:, 2], box1[:, 3] - b2_x1, b2_y1, b2_x2, b2_y2 = box2[:, 0], box2[:, 1], box2[:, 2], box2[:, 3] +def post_process( + boxes: np.ndarray, + scores: np.ndarray, + classes: np.ndarray, + *, + conf_threshold: float, + iou_threshold: float, + origin_h: int, + origin_w: int, +) -> list[Detection]: + """Filter by confidence, run NMS, return a :class:`Detection` list in x/y/w/h form.""" + mask = scores > conf_threshold + boxes = boxes[mask].copy() + scores = scores[mask] + classes = classes[mask] - inter_x1 = np.maximum(b1_x1, b2_x1) - inter_y1 = np.maximum(b1_y1, b2_y1) - inter_x2 = np.minimum(b1_x2, b2_x2) - inter_y2 = np.minimum(b1_y2, b2_y2) + if len(scores) == 0: + return [] - inter_area = np.clip(inter_x2 - inter_x1 + 1, 0, None) * np.clip(inter_y2 - inter_y1 + 1, 0, None) - b1_area = (b1_x2 - b1_x1 + 1) * (b1_y2 - b1_y1 + 1) - b2_area = (b2_x2 - b2_x1 + 1) * (b2_y2 - b2_y1 + 1) + boxes, scores, classes = non_max_suppression( + boxes, scores, classes, + iou_threshold=iou_threshold, origin_h=origin_h, origin_w=origin_w) - return inter_area / (b1_area + b2_area - inter_area + 1e-16) + result = [] + for j, box in enumerate(boxes): + x1, y1, x2, y2 = box + w = x2 - x1 + h = y2 - y1 + result.append(Detection(int(x1), int(y1), int(w), int(h), + int(classes[j]), round(float(scores[j]), 2))) + return result def non_max_suppression( @@ -99,37 +112,24 @@ def non_max_suppression( return boxes[keep_indices], scores[keep_indices], classes[keep_indices] -def post_process( - boxes: np.ndarray, - scores: np.ndarray, - classes: np.ndarray, - *, - conf_threshold: float, - iou_threshold: float, - origin_h: int, - origin_w: int, -) -> list[Detection]: - """Filter by confidence, run NMS, return a :class:`Detection` list in x/y/w/h form.""" - mask = scores > conf_threshold - boxes = boxes[mask].copy() - scores = scores[mask] - classes = classes[mask] +def bbox_iou( + box1: np.ndarray, + box2: np.ndarray, +) -> np.ndarray: + """Compute IoU between box1 (1x4) and box2 (Nx4), both in x1y1x2y2 format.""" + b1_x1, b1_y1, b1_x2, b1_y2 = box1[:, 0], box1[:, 1], box1[:, 2], box1[:, 3] + b2_x1, b2_y1, b2_x2, b2_y2 = box2[:, 0], box2[:, 1], box2[:, 2], box2[:, 3] - if len(scores) == 0: - return [] + inter_x1 = np.maximum(b1_x1, b2_x1) + inter_y1 = np.maximum(b1_y1, b2_y1) + inter_x2 = np.minimum(b1_x2, b2_x2) + inter_y2 = np.minimum(b1_y2, b2_y2) - boxes, scores, classes = non_max_suppression( - boxes, scores, classes, - iou_threshold=iou_threshold, origin_h=origin_h, origin_w=origin_w) + inter_area = np.clip(inter_x2 - inter_x1 + 1, 0, None) * np.clip(inter_y2 - inter_y1 + 1, 0, None) + b1_area = (b1_x2 - b1_x1 + 1) * (b1_y2 - b1_y1 + 1) + b2_area = (b2_x2 - b2_x1 + 1) * (b2_y2 - b2_y1 + 1) - result = [] - for j, box in enumerate(boxes): - x1, y1, x2, y2 = box - w = x2 - x1 - h = y2 - y1 - result.append(Detection(int(x1), int(y1), int(w), int(h), - int(classes[j]), round(float(scores[j]), 2))) - return result + return inter_area / (b1_area + b2_area - inter_area + 1e-16) def detections_from_xyxy( @@ -152,6 +152,32 @@ def detections_from_xyxy( return result +def to_image_metadata( + detections: list[Detection], + model_information: ModelInformation, + im_height: int, + im_width: int, +) -> ImageMetadata: + """Build the container a *detector* node reports from a list of detections.""" + image_metadata = ImageMetadata() + _append_detections(image_metadata, detections, model_information, im_height, im_width) + return image_metadata + + +def to_detections( + detections: list[Detection], + model_information: ModelInformation, + im_height: int, + im_width: int, + *, + image_id: str | None = None, +) -> Detections: + """Build the container a *trainer*'s auto-detection pass reports.""" + result = Detections(image_id=image_id) + _append_detections(result, detections, model_information, im_height, im_width) + return result + + def _append_detections( target: ImageMetadata | Detections, detections: list[Detection], @@ -208,29 +234,3 @@ def _append_detections( if skipped_detections: log_msg = '\n'.join([str(d) for d in skipped_detections]) logging.warning('Removed %d small detections from result: \n%s', len(skipped_detections), log_msg) - - -def to_image_metadata( - detections: list[Detection], - model_information: ModelInformation, - im_height: int, - im_width: int, -) -> ImageMetadata: - """Build the container a *detector* node reports from a list of detections.""" - image_metadata = ImageMetadata() - _append_detections(image_metadata, detections, model_information, im_height, im_width) - return image_metadata - - -def to_detections( - detections: list[Detection], - model_information: ModelInformation, - im_height: int, - im_width: int, - *, - image_id: str | None = None, -) -> Detections: - """Build the container a *trainer*'s auto-detection pass reports.""" - result = Detections(image_id=image_id) - _append_detections(result, detections, model_information, im_height, im_width) - return result diff --git a/learning_loop_node/trainer/subprocess.py b/learning_loop_node/trainer/subprocess.py index dbfa181b..45afb4ed 100644 --- a/learning_loop_node/trainer/subprocess.py +++ b/learning_loop_node/trainer/subprocess.py @@ -25,24 +25,17 @@ P = ParamSpec('P') -class IteratorDone: - pass - - -def _iterator_wrapper( - it: Callable[..., Iterator[T]], - state_queue: MPQueue[T | Exception | IteratorDone], - args: tuple[Any, ...], - kwargs: dict[str, Any], -) -> None: +@asynccontextmanager +async def iterator_cpu_bound( + it: Callable[P, Iterator[T]], + *args: P.args, + **kwargs: P.kwargs, +) -> AsyncGenerator[AsyncGenerator[T, None], None]: + iterator = _iterator_cpu_bound_inner(it, *args, **kwargs) try: - for data in it(*args, **kwargs): - state_queue.put(data) - except Exception as e: - print(traceback.format_exc()) - state_queue.put(e) - - state_queue.put(IteratorDone()) + yield iterator + finally: + await asyncio.shield(iterator.aclose()) async def _iterator_cpu_bound_inner( @@ -80,14 +73,21 @@ async def _iterator_cpu_bound_inner( process.join() -@asynccontextmanager -async def iterator_cpu_bound( - it: Callable[P, Iterator[T]], - *args: P.args, - **kwargs: P.kwargs, -) -> AsyncGenerator[AsyncGenerator[T, None], None]: - iterator = _iterator_cpu_bound_inner(it, *args, **kwargs) +def _iterator_wrapper( + it: Callable[..., Iterator[T]], + state_queue: MPQueue[T | Exception | IteratorDone], + args: tuple[Any, ...], + kwargs: dict[str, Any], +) -> None: try: - yield iterator - finally: - await asyncio.shield(iterator.aclose()) + for data in it(*args, **kwargs): + state_queue.put(data) + except Exception as e: + print(traceback.format_exc()) + state_queue.put(e) + + state_queue.put(IteratorDone()) + + +class IteratorDone: + pass From dd582cefbdba9ccb599a628e9fbad8b144afda2c Mon Sep 17 00:00:00 2001 From: Kevin Heye Date: Wed, 26 Aug 2026 15:22:00 +0200 Subject: [PATCH 12/29] Log through the module logger, not the root one CONTRIBUTING rule 10 asks for a module-level logger. The three modules this branch adds or changes were logging through the root logger, inherited from the node code they came from: `postprocess.py` twice, `entrypoint.py` twice and `environment_reader.py` twice, including the warning this branch added. `getLogger(__name__)` rather than a hand-written name, matching `batch_size.py` and the newer library modules. Also modernise `read_from_env`'s annotations while the file is open, which is what rule 7 asks for and clears three pre-existing ruff findings (`List`, `Optional`). 203 unit tests pass, ruff clean on all four files. Co-Authored-By: Claude Opus 5 --- learning_loop_node/detector/postprocess.py | 6 ++++-- learning_loop_node/helpers/entrypoint.py | 6 ++++-- learning_loop_node/helpers/environment_reader.py | 9 +++++---- 3 files changed, 13 insertions(+), 8 deletions(-) diff --git a/learning_loop_node/detector/postprocess.py b/learning_loop_node/detector/postprocess.py index 0fb509fa..115c126e 100644 --- a/learning_loop_node/detector/postprocess.py +++ b/learning_loop_node/detector/postprocess.py @@ -28,6 +28,8 @@ from .categories import category_by_index from .geometry import clip_box, clip_point +logger = logging.getLogger(__name__) + MIN_BOX_SIZE: int = 2 """Boxes this small are dropped: they carry no usable information and clutter the loop.""" @@ -229,8 +231,8 @@ def _append_detections( ) ) else: - logging.warning('Unsupported category type %s for category %s', category.type, category.name) + logger.warning('Unsupported category type %s for category %s', category.type, category.name) if skipped_detections: log_msg = '\n'.join([str(d) for d in skipped_detections]) - logging.warning('Removed %d small detections from result: \n%s', len(skipped_detections), log_msg) + logger.warning('Removed %d small detections from result: \n%s', len(skipped_detections), log_msg) diff --git a/learning_loop_node/helpers/entrypoint.py b/learning_loop_node/helpers/entrypoint.py index 1b360c76..46dad3ca 100644 --- a/learning_loop_node/helpers/entrypoint.py +++ b/learning_loop_node/helpers/entrypoint.py @@ -20,6 +20,8 @@ import configargparse import uvicorn +logger = logging.getLogger(__name__) + def node_parser(*, description: str, legacy_env_prefix: str = '') -> configargparse.ArgumentParser: """Build the parser for a node, pre-loaded with the settings every node has. @@ -42,7 +44,7 @@ def run_node(app: str, args: Namespace) -> None: :param app: Import string of the node object, conventionally ``'main:node'``. """ reload = os.getenv('UVICORN_RELOAD', 'FALSE').lower() in ('true', '1') - logging.info('Uvicorn reload is set to: %s', reload) + logger.info('Uvicorn reload is set to: %s', reload) uvicorn.run(app, host=args.host, port=args.port, lifespan='on', reload=reload) @@ -72,4 +74,4 @@ def _adopt_legacy_env_vars(self) -> None: if not legacy or name in os.environ or legacy not in os.environ: continue os.environ[name] = os.environ[legacy] - logging.warning('%s is deprecated and will stop being read; set %s instead', legacy, name) + logger.warning('%s is deprecated and will stop being read; set %s instead', legacy, name) diff --git a/learning_loop_node/helpers/environment_reader.py b/learning_loop_node/helpers/environment_reader.py index 3cb8cca3..84d77ac6 100644 --- a/learning_loop_node/helpers/environment_reader.py +++ b/learning_loop_node/helpers/environment_reader.py @@ -1,17 +1,18 @@ import logging import os -from typing import List, Optional + +logger = logging.getLogger(__name__) # TODO ignore_errors should default to False, but maybe some tests rely on this behavior -def read_from_env(possible_names: List[str], ignore_errors: bool = True) -> Optional[str]: +def read_from_env(possible_names: list[str], ignore_errors: bool = True) -> str | None: values = [os.environ.get(name, None) for name in possible_names] values = list(filter(None, values)) # Possible error: no values are set if not values: if ignore_errors: - logging.warning('no environment variable set for %s', possible_names) + logger.warning('no environment variable set for %s', possible_names) return None raise ValueError(f'no environment variable set for {possible_names}') @@ -22,7 +23,7 @@ def read_from_env(possible_names: List[str], ignore_errors: bool = True) -> Opti if len(values) > 1 and len(set(values)) > 1: if not ignore_errors: raise ValueError(f'different environment variables set for {possible_names}: {values}') - logging.warning('different environment variables set for %s: %s - using %s', + logger.warning('different environment variables set for %s: %s - using %s', possible_names, values, values[0]) return values[0] From c638660de62a379c3499aa502af5eab8c3c20876 Mon Sep 17 00:00:00 2001 From: Kevin Heye Date: Wed, 26 Aug 2026 16:17:24 +0200 Subject: [PATCH 13/29] Read a renamed setting under the name its prefix actually gave it A node that used to require an environment prefix had it applied to the *flag*, so `--host` was read from `HOST`. `_adopt_legacy_env_vars` derived the legacy name from the current `env_var` instead, so for `--host` it looked for `NODE_HOST` -- a name no deployment has ever set, because the rename to `NODE_HOST` is what this migration introduces. The effect was worst for exactly the settings the migration renamed: `HOST` and `PORT` were dropped silently, with no value adopted and no deprecation warning, and the node fell back to binding 0.0.0.0:80. Every setting whose environment variable is simply its flag in upper case was unaffected, which is why it went unnoticed. Both spellings are now tried, the prefixed current name first, so a node whose prefixed spelling already matched keeps working too. Co-Authored-By: Claude Opus 5 --- learning_loop_node/helpers/entrypoint.py | 28 ++++++++++--- .../tests/unit/test_entrypoint.py | 40 ++++++++++++++++++- 2 files changed, 61 insertions(+), 7 deletions(-) diff --git a/learning_loop_node/helpers/entrypoint.py b/learning_loop_node/helpers/entrypoint.py index 46dad3ca..4ec7dbe8 100644 --- a/learning_loop_node/helpers/entrypoint.py +++ b/learning_loop_node/helpers/entrypoint.py @@ -28,8 +28,10 @@ def node_parser(*, description: str, legacy_env_prefix: str = '') -> configargpa :param legacy_env_prefix: A prefix an earlier version of this node required, e.g. ``'MY_DETECTOR_'``. Prefixed names are still honoured, with a warning, so a - deployment keeps working until it is updated. Leave empty for a node that has always - read unprefixed names. + deployment keeps working until it is updated. Both spellings are accepted: the prefix + on the current environment variable (``MY_DETECTOR_NODE_HOST``) and the prefix on the + flag it was originally applied to (``MY_DETECTOR_HOST``). Leave empty for a node that + has always read unprefixed names. """ parser = _NodeArgumentParser(description=description, legacy_env_prefix=legacy_env_prefix) parser.add_argument('--host', default='0.0.0.0', env_var='NODE_HOST', @@ -70,8 +72,22 @@ def _adopt_legacy_env_vars(self) -> None: return for action in self._actions: name = getattr(action, 'env_var', None) - legacy = self.legacy_env_prefix + name if name else None - if not legacy or name in os.environ or legacy not in os.environ: + if not name or name in os.environ: continue - os.environ[name] = os.environ[legacy] - logger.warning('%s is deprecated and will stop being read; set %s instead', legacy, name) + for legacy in self._legacy_names(action, name): + if legacy not in os.environ: + continue + os.environ[name] = os.environ[legacy] + logger.warning('%s is deprecated and will stop being read; set %s instead', legacy, name) + break + + def _legacy_names(self, action: Action, name: str) -> list[str]: + """The prefixed names this setting may still be configured under, preferred first. + + A prefix used to be applied to the *flag*, so the old name of ``--host`` was + ``HOST`` — not ``NODE_HOST``. Deriving the legacy name from the current + ``env_var`` alone therefore misses exactly the settings that were renamed, which is every + one whose environment variable is not simply its flag in upper case. + """ + candidates = [self.legacy_env_prefix + name, self.legacy_env_prefix + action.dest.upper()] + return list(dict.fromkeys(candidates)) diff --git a/learning_loop_node/tests/unit/test_entrypoint.py b/learning_loop_node/tests/unit/test_entrypoint.py index d20a0911..03a29f6b 100644 --- a/learning_loop_node/tests/unit/test_entrypoint.py +++ b/learning_loop_node/tests/unit/test_entrypoint.py @@ -2,7 +2,8 @@ from ...helpers.entrypoint import node_parser -MANAGED = ('WEIGHT_TYPE', 'MY_DETECTOR_WEIGHT_TYPE', 'HOST', 'NODE_HOST', 'NODE_PORT', 'PORT') +MANAGED = ('WEIGHT_TYPE', 'MY_DETECTOR_WEIGHT_TYPE', 'HOST', 'NODE_HOST', 'NODE_PORT', 'PORT', + 'MY_DETECTOR_HOST', 'MY_DETECTOR_PORT', 'MY_DETECTOR_NODE_HOST') @pytest.fixture(autouse=True) @@ -69,3 +70,40 @@ def test_the_current_name_wins_over_the_prefixed_one(monkeypatch: pytest.MonkeyP def test_a_node_without_a_legacy_prefix_ignores_prefixed_names(monkeypatch: pytest.MonkeyPatch): monkeypatch.setenv('MY_DETECTOR_WEIGHT_TYPE', 'FP32') assert _parser().parse_args([]).weight_type == 'FP16' + + +def test_the_bind_address_is_still_read_under_the_name_the_prefix_gave_the_flag( + monkeypatch: pytest.MonkeyPatch): + """A prefix used to be applied to the flag, so `--host` was `HOST`, not + `NODE_HOST`. That is the only spelling a deployment can have set.""" + monkeypatch.setenv('MY_DETECTOR_HOST', '127.0.0.1') + monkeypatch.setenv('MY_DETECTOR_PORT', '8099') + args = _parser(legacy_env_prefix='MY_DETECTOR_').parse_args([]) + assert (args.host, args.port) == ('127.0.0.1', 8099) + + +def test_the_renamed_setting_warns_which_name_to_use_instead(monkeypatch: pytest.MonkeyPatch, + caplog: pytest.LogCaptureFixture): + monkeypatch.setenv('MY_DETECTOR_HOST', '127.0.0.1') + _parser(legacy_env_prefix='MY_DETECTOR_').parse_args([]) + assert 'MY_DETECTOR_HOST' in caplog.text + assert 'NODE_HOST' in caplog.text + + +def test_the_prefixed_current_name_is_honoured_too(monkeypatch: pytest.MonkeyPatch): + """A node whose prefixed spelling was already the current one keeps working.""" + monkeypatch.setenv('MY_DETECTOR_NODE_HOST', '10.0.0.1') + assert _parser(legacy_env_prefix='MY_DETECTOR_').parse_args([]).host == '10.0.0.1' + + +def test_the_current_name_wins_over_both_prefixed_spellings(monkeypatch: pytest.MonkeyPatch): + monkeypatch.setenv('NODE_HOST', '10.0.0.1') + monkeypatch.setenv('MY_DETECTOR_NODE_HOST', '127.0.0.2') + monkeypatch.setenv('MY_DETECTOR_HOST', '127.0.0.3') + assert _parser(legacy_env_prefix='MY_DETECTOR_').parse_args([]).host == '10.0.0.1' + + +def test_the_loop_own_host_is_not_adopted_by_a_prefixed_node(monkeypatch: pytest.MonkeyPatch): + """`HOST` is the loop's address even here: only the *prefixed* spelling is a bind address.""" + monkeypatch.setenv('HOST', 'preview.learning-loop.ai') + assert _parser(legacy_env_prefix='MY_DETECTOR_').parse_args([]).host == '0.0.0.0' From 943959a59ee5bf190c6a2fe656f2339473337620 Mon Sep 17 00:00:00 2001 From: Kevin Heye Date: Wed, 2 Sep 2026 14:09:56 +0200 Subject: [PATCH 14/29] Carry main's improvements to the moved files into the library copies The node repositories improved these two files after this branch had already moved them out, so merging main there resolves towards the deletion. The improvements themselves still belong here, or the move would quietly undo them: - `Detection` becomes a typed `NamedTuple`. As a bare `namedtuple` its six fields had no types at all, which is what dfine_node fixed in #39. - The failing child process logs through the logger instead of `print(traceback.format_exc())`, so a training failure lands in the node's log rather than only on stdout. The module logger, not the root one, since that is what the rest of this library does. Assisted-by: Claude:claude-opus-5 Co-Authored-By: Claude Opus 5 --- learning_loop_node/detector/postprocess.py | 16 ++++++++++++---- learning_loop_node/trainer/subprocess.py | 6 ++++-- 2 files changed, 16 insertions(+), 6 deletions(-) diff --git a/learning_loop_node/detector/postprocess.py b/learning_loop_node/detector/postprocess.py index 115c126e..8cffc708 100644 --- a/learning_loop_node/detector/postprocess.py +++ b/learning_loop_node/detector/postprocess.py @@ -13,7 +13,7 @@ """ import logging -from collections import namedtuple +from typing import NamedTuple import numpy as np @@ -33,9 +33,17 @@ MIN_BOX_SIZE: int = 2 """Boxes this small are dropped: they carry no usable information and clutter the loop.""" -Detection = namedtuple('Detection', 'x y w h category probability') -"""One surviving prediction. ``x``/``y`` are the top-left corner, ``category`` is an index -into :attr:`ModelInformation.categories`.""" + +class Detection(NamedTuple): + """One surviving prediction. ``x``/``y`` are the top-left corner, ``category`` is an index + into :attr:`ModelInformation.categories`.""" + + x: int + y: int + w: int + h: int + category: int + probability: float def post_process( diff --git a/learning_loop_node/trainer/subprocess.py b/learning_loop_node/trainer/subprocess.py index 45afb4ed..ae6e178c 100644 --- a/learning_loop_node/trainer/subprocess.py +++ b/learning_loop_node/trainer/subprocess.py @@ -13,14 +13,16 @@ from __future__ import annotations import asyncio +import logging import multiprocessing import queue -import traceback from collections.abc import AsyncGenerator, Callable, Iterator from contextlib import asynccontextmanager from multiprocessing.queues import Queue as MPQueue from typing import Any, ParamSpec, TypeVar +logger = logging.getLogger(__name__) + T = TypeVar('T') P = ParamSpec('P') @@ -83,7 +85,7 @@ def _iterator_wrapper( for data in it(*args, **kwargs): state_queue.put(data) except Exception as e: - print(traceback.format_exc()) + logger.exception('iterator_cpu_bound child process failed') state_queue.put(e) state_queue.put(IteratorDone()) From 8641a8a14921f03bc8915068a178158accc5fa4d Mon Sep 17 00:00:00 2001 From: jan Date: Thu, 3 Sep 2026 11:51:06 +0200 Subject: [PATCH 15/29] Apply the conventions review to the code this branch adds The docstrings written while moving this code explain why each implementation was chosen, which CONTRIBUTING rule 4 puts in the pull request rather than in code that goes stale after the merge; what the code cannot show stays, and the NOTE block in read_from_env is reduced to its one real point, that possible_names is ordered by preference. find_batch_size now raises InsufficientMemoryError instead of the bare RuntimeError that is_out_of_memory exists to disambiguate, Detection's fields are named at the call site (rule 9), and the docstrings in trainer_logic.py and exceptions.py use double quotes (rule 1). subprocess.py also drops the claim that the process is "spawned" - without an explicit context multiprocessing takes the platform default, fork on Linux. --- learning_loop_node/detector/categories.py | 8 +----- learning_loop_node/detector/geometry.py | 10 +++----- learning_loop_node/detector/postprocess.py | 19 ++++++-------- learning_loop_node/helpers/entrypoint.py | 16 ++++-------- .../helpers/environment_reader.py | 12 +++++---- learning_loop_node/trainer/batch_size.py | 25 ++++++++----------- learning_loop_node/trainer/exceptions.py | 12 ++++++--- learning_loop_node/trainer/metrics.py | 6 ++--- learning_loop_node/trainer/subprocess.py | 9 +++---- learning_loop_node/trainer/trainer_logic.py | 19 +++++++------- 10 files changed, 59 insertions(+), 77 deletions(-) diff --git a/learning_loop_node/detector/categories.py b/learning_loop_node/detector/categories.py index 9f45541f..a697ce83 100644 --- a/learning_loop_node/detector/categories.py +++ b/learning_loop_node/detector/categories.py @@ -1,9 +1,7 @@ """Resolving a model's categories from what it emits. A model reports class indices, or names, and only ``ModelInformation.categories`` gives those -meaning. Both lookups are the point at which a model and its metadata are checked against each -other, so both raise rather than skipping quietly: a mismatch here makes every prediction on the -image suspect, and it is far cheaper to diagnose at the lookup than three layers later. +meaning. """ from ..data_classes import Category, ModelInformation @@ -12,10 +10,6 @@ def category_by_index(model_information: ModelInformation, index: int) -> Category: """Resolve the category a model's class index refers to. - Models emit class indices, and the order of ``model_information.categories`` is what - gives them meaning — so an out-of-range index is a model/metadata mismatch, not a - detection to skip quietly. - :raises ValueError: If the index is outside the model's category list. """ categories = model_information.categories diff --git a/learning_loop_node/detector/geometry.py b/learning_loop_node/detector/geometry.py index b17ed33b..db81abdc 100644 --- a/learning_loop_node/detector/geometry.py +++ b/learning_loop_node/detector/geometry.py @@ -1,10 +1,8 @@ """Box and point clipping shared by every detector node. -The loop stores a box as its top-left corner plus a size, so :func:`clip_box` is the form a -node needs when it hands detections to the loop. Model outputs are not always in that form — -:func:`clip_box_centered` keeps the centre-based convention explicit instead of letting two -incompatible functions share one name, which is how the same helper ended up meaning two -different things in different node repositories. +The loop stores a box as its top-left corner plus a size, which is the form :func:`clip_box` +produces. :func:`clip_box_centered` is the centre-anchored variant, for model outputs in that +convention. """ @@ -21,7 +19,7 @@ def clip_box( :param x1: Left edge of the box. :param y1: Top edge of the box. - :return: The clipped ``(x1, y1, width, height)`` as ints; the size is never negative. + :return: The clipped ``(x1, y1, width, height)``; the size is never negative. """ x2 = x1 + width y2 = y1 + height diff --git a/learning_loop_node/detector/postprocess.py b/learning_loop_node/detector/postprocess.py index 8cffc708..e544197b 100644 --- a/learning_loop_node/detector/postprocess.py +++ b/learning_loop_node/detector/postprocess.py @@ -2,14 +2,12 @@ Every detection node ends up doing the same three things with a model's raw output: drop low-confidence predictions, suppress overlapping boxes, and turn what survives into the -loop's detection dataclasses. None of that depends on the model, so it lives here rather -than being re-derived — and re-diverging — in each node repository. +loop's detection dataclasses. None of that depends on the model. Two containers carry the same detections in this library: a detector node reports :class:`~learning_loop_node.data_classes.image_metadata.ImageMetadata`, while a trainer's auto-detection pass reports :class:`~learning_loop_node.data_classes.detections.Detections`. -:func:`to_image_metadata` and :func:`to_detections` build them from the same routine, so both -paths clip and filter identically. +:func:`to_image_metadata` and :func:`to_detections` build them from the same routine. """ import logging @@ -31,7 +29,7 @@ logger = logging.getLogger(__name__) MIN_BOX_SIZE: int = 2 -"""Boxes this small are dropped: they carry no usable information and clutter the loop.""" +"""A box is dropped unless both of its sides exceed this.""" class Detection(NamedTuple): @@ -74,8 +72,8 @@ def post_process( x1, y1, x2, y2 = box w = x2 - x1 h = y2 - y1 - result.append(Detection(int(x1), int(y1), int(w), int(h), - int(classes[j]), round(float(scores[j]), 2))) + result.append(Detection(x=int(x1), y=int(y1), w=int(w), h=int(h), + category=int(classes[j]), probability=round(float(scores[j]), 2))) return result @@ -151,14 +149,13 @@ def detections_from_xyxy( """Convert already-suppressed model output into :class:`Detection` values. For nodes whose model (or a torch/ONNX op) has done the suppression already, so only the - coordinate conversion is left. Corners are rounded rather than truncated, which is half a - pixel more faithful than :func:`post_process` — that one keeps truncating so its output - stays bit-identical to what detectors reported before this module existed. + coordinate conversion is left. Corners are rounded here; :func:`post_process` truncates. """ result = [] for label, box, score in zip(labels, boxes, scores, strict=True): x1, y1, x2, y2 = (round(value) for value in box) - result.append(Detection(x1, y1, x2 - x1, y2 - y1, int(label), score)) + result.append(Detection(x=x1, y=y1, w=x2 - x1, h=y2 - y1, + category=int(label), probability=score)) return result diff --git a/learning_loop_node/helpers/entrypoint.py b/learning_loop_node/helpers/entrypoint.py index 4ec7dbe8..5f0d6898 100644 --- a/learning_loop_node/helpers/entrypoint.py +++ b/learning_loop_node/helpers/entrypoint.py @@ -4,13 +4,9 @@ logic object, construct the node, and hand it to uvicorn. Only the middle two are the node's own, so :func:`node_parser` and :func:`run_node` cover the rest. -Settings come from a flag *or* an environment variable, because a node is configured on the -command line while developing and through the container environment in deployment. One -declaration gives both: ``--conf-threshold`` reads ``CONF_THRESHOLD``. - -``--host`` and ``--port`` are the exception — they read ``NODE_HOST`` and ``NODE_PORT`` rather -than the names their flags imply, because the bare ``HOST`` already means *the loop's address* -and a node adopting it would hand it to uvicorn and fail to bind. +Every setting is a flag *and* an environment variable named after it: ``--conf-threshold`` +reads ``CONF_THRESHOLD``. ``--host`` and ``--port`` are the exception, reading ``NODE_HOST`` +and ``NODE_PORT``: the bare ``HOST`` is the address of the loop. """ import logging @@ -84,10 +80,8 @@ def _adopt_legacy_env_vars(self) -> None: def _legacy_names(self, action: Action, name: str) -> list[str]: """The prefixed names this setting may still be configured under, preferred first. - A prefix used to be applied to the *flag*, so the old name of ``--host`` was - ``HOST`` — not ``NODE_HOST``. Deriving the legacy name from the current - ``env_var`` alone therefore misses exactly the settings that were renamed, which is every - one whose environment variable is not simply its flag in upper case. + A prefix used to be applied to the *flag*, so the old name of ``--host`` is + ``HOST``, not ``NODE_HOST``. """ candidates = [self.legacy_env_prefix + name, self.legacy_env_prefix + action.dest.upper()] return list(dict.fromkeys(candidates)) diff --git a/learning_loop_node/helpers/environment_reader.py b/learning_loop_node/helpers/environment_reader.py index 84d77ac6..002652f0 100644 --- a/learning_loop_node/helpers/environment_reader.py +++ b/learning_loop_node/helpers/environment_reader.py @@ -6,6 +6,11 @@ # TODO ignore_errors should default to False, but maybe some tests rely on this behavior def read_from_env(possible_names: list[str], ignore_errors: bool = True) -> str | None: + """Read the first of ``possible_names`` that is set. + + :param possible_names: In order of preference; on a disagreement the first one set wins. + :raises ValueError: If nothing is set or the values disagree, unless ``ignore_errors``. + """ values = [os.environ.get(name, None) for name in possible_names] values = list(filter(None, values)) @@ -16,15 +21,12 @@ def read_from_env(possible_names: list[str], ignore_errors: bool = True) -> str return None raise ValueError(f'no environment variable set for {possible_names}') - # Possible error: multiple values are not None and not equal. - # NOTE returning None here would be worse than picking one: the caller falls back to its - # default, and host()'s default is the production loop. `possible_names` is ordered by - # preference, so a disagreement resolves to the first name that is set. + # Possible error: multiple values are not None and not equal if len(values) > 1 and len(set(values)) > 1: if not ignore_errors: raise ValueError(f'different environment variables set for {possible_names}: {values}') logger.warning('different environment variables set for %s: %s - using %s', - possible_names, values, values[0]) + possible_names, values, values[0]) return values[0] diff --git a/learning_loop_node/trainer/batch_size.py b/learning_loop_node/trainer/batch_size.py index 707bb8c6..0e438036 100644 --- a/learning_loop_node/trainer/batch_size.py +++ b/learning_loop_node/trainer/batch_size.py @@ -1,16 +1,11 @@ """Choosing a batch size by probing, rather than configuring one. -The largest batch that fits depends on the model, the image resolution and the card, so a -configured value is either wrong on some machines or too small on all of them. Probing answers -it directly: run a representative step at doubling sizes and keep the last one that survived. +The largest batch that fits depends on the model, the image resolution and the card. Probing +runs a representative step at doubling sizes and keeps the last one that survived; what "a +representative step" means is the trainer's business — it supplies a ``fits`` predicate. +Nothing here imports a deep-learning framework. -What "a representative step" means is the trainer's business — it supplies a ``fits`` predicate -that runs a real step and reports whether it ran out of memory. Everything here is that -predicate's scaffolding, and none of it imports a deep-learning framework. - -Only powers of two are visited, so equal hardware yields an equal recipe. That matters when a -trainer scales its learning rate by the batch size: a probe that returned 37 on one machine and -39 on another would make the two trainings quietly different. +Only powers of two are visited, so equal hardware yields an equal recipe. Adapted from PyTorch Lightning's ``BatchSizeFinder`` (power-scaling mode). Copyright The Lightning AI team. Licensed under the Apache License, Version 2.0. @@ -20,24 +15,26 @@ import logging from collections.abc import Callable +from .exceptions import InsufficientMemoryError + logger = logging.getLogger(__name__) NO_GPU_BATCH_SIZE = 8 -"""What to fall back on with no GPU to probe: enough to make progress, small enough to fit.""" +"""Batch size used when there is no GPU to probe.""" MIN_TRAIN_STEPS_PER_EPOCH = 8 -"""Below this an epoch is one or two optimizer steps, which trains poorly whatever the card.""" +"""Fewest optimizer steps an epoch must have; :func:`dataset_limit` is derived from it.""" def find_batch_size(fits: Callable[[int], bool], *, limit: int) -> int: """Return the largest power-of-two batch size that fits, never exceeding ``limit``. :param fits: Runs a representative probe; ``False`` on out-of-memory. - :raises RuntimeError: If not even a batch size of 1 fits. + :raises InsufficientMemoryError: If not even a batch size of 1 fits. """ limit = smaller_pot(limit) if not fits(1): - raise RuntimeError('batch size 1 does not fit in memory') + raise InsufficientMemoryError('batch size 1 does not fit in memory') size = 1 while size < limit and fits(size * 2): diff --git a/learning_loop_node/trainer/exceptions.py b/learning_loop_node/trainer/exceptions.py index c3cb54f2..b77548cd 100644 --- a/learning_loop_node/trainer/exceptions.py +++ b/learning_loop_node/trainer/exceptions.py @@ -1,12 +1,16 @@ class CriticalError(Exception): - ''' + """ CriticalError is raised when the training cannot be continued. In this case the trainer jumps to the TrainerState.ReadyForCleanup and tries to upload the latest model. - ''' + """ + + +class InsufficientMemoryError(RuntimeError): + """Raised when not even the smallest unit of work fits in memory.""" class NodeNeedsRestartError(Exception): - ''' + """ NodeNeedsRestartError is raised when the node needs to be restarted. This is e.g. the case when the GPU is not available anymore. - ''' + """ diff --git a/learning_loop_node/trainer/metrics.py b/learning_loop_node/trainer/metrics.py index aee553b8..6fb6baf9 100644 --- a/learning_loop_node/trainer/metrics.py +++ b/learning_loop_node/trainer/metrics.py @@ -1,8 +1,7 @@ """Scoring a training from the confusion matrix the loop stores. The loop keeps one ``{'tp': .., 'fp': .., 'fn': ..}`` per category, which is what -``TrainerLogicGeneric._get_new_best_training_state`` returns. Deciding whether an epoch beat -the previous best means reducing that to one number, and every trainer needs the same one. +``TrainerLogicGeneric._get_new_best_training_state`` returns. """ import statistics @@ -12,8 +11,7 @@ def macro_f1(confusion_matrix: Mapping[str, Mapping[str, int]]) -> float: """The unweighted mean of the per-category F1 scores, as the loop's UI shows it by default. - Unweighted is the point: a rare category the model never finds drags the score down as much - as a frequent one, where pooling the counts first would hide it. + Unweighted, so a rare category weighs as much as a frequent one. """ scores = [category_f1(counts) for counts in confusion_matrix.values()] return statistics.mean(scores) if scores else 0.0 diff --git a/learning_loop_node/trainer/subprocess.py b/learning_loop_node/trainer/subprocess.py index ae6e178c..4f038547 100644 --- a/learning_loop_node/trainer/subprocess.py +++ b/learning_loop_node/trainer/subprocess.py @@ -1,11 +1,8 @@ """Run a blocking, CPU-bound generator in its own process without blocking the event loop. -A trainer that trains in-process has a problem: the training must not stall the node, and CUDA -state must stay out of the node process so a crashed training cannot take the node with it. -:func:`iterator_cpu_bound` runs the generator in a spawned process and yields what it produces -through a ``maxsize=1`` queue, so the producer can never run more than one item ahead of the -bookkeeping that consumes it — which is what lets a trainer alternate between two model files -and know the one it is copying is not being rewritten. +:func:`iterator_cpu_bound` runs the generator in a separate process and yields what it produces +through a ``maxsize=1`` queue, so the producer never runs more than one item ahead of the +consumer. Exceptions raised inside the process are re-raised in the caller, and the process is killed if the caller leaves the context early. diff --git a/learning_loop_node/trainer/trainer_logic.py b/learning_loop_node/trainer/trainer_logic.py index 9827b41b..b786113f 100644 --- a/learning_loop_node/trainer/trainer_logic.py +++ b/learning_loop_node/trainer/trainer_logic.py @@ -161,28 +161,29 @@ async def stop(self) -> None: @abstractmethod async def _start_training_from_base_model(self) -> None: - '''Should be used to start a training on executer, e.g. self.executor.start(cmd).''' + """Should be used to start a training on executer, e.g. self.executor.start(cmd).""" @abstractmethod async def _start_training_from_scratch(self) -> None: - '''Should be used to start a training from scratch on executer, e.g. self.executor.start(cmd). + """Should be used to start a training from scratch on executer, e.g. self.executor.start(cmd). NOTE base_model_id is now accessible via self.training.base_model_id - the id of a pretrained model provided by self.provided_pretrained_models.''' + the id of a pretrained model provided by self.provided_pretrained_models.""" @abstractmethod def _can_resume(self) -> bool: - '''Override this method to return True if the trainer can resume training.''' + """Override this method to return True if the trainer can resume training.""" @abstractmethod async def _resume(self) -> None: - '''Is called when self.can_resume() returns True. - One may resume the training on a previously trained model stored by self.on_model_published(basic_model).''' + """Is called when self.can_resume() returns True. + One may resume the training on a previously trained model stored by self.on_model_published(basic_model).""" def _get_executor_error_from_log(self) -> Optional[str]: - '''Reports what went wrong to the Learning Loop by reading self.executor's log. + """Reports what went wrong to the Learning Loop by reading self.executor's log. The default recognises the CUDA failures every trainer hits. Override to add messages a - particular training framework produces, and call super() to keep these.''' + particular training framework produces, and call super() to keep these. + """ if self._executor is None: return None for line in self._executor.get_log_by_lines(tail=50): @@ -194,4 +195,4 @@ def _get_executor_error_from_log(self) -> Optional[str]: @abstractmethod async def _detect(self, model_information: ModelInformation, images: List[str], model_folder: str) -> List[Detections]: - '''Called to run detections on a list of images.''' + """Called to run detections on a list of images.""" From 915bedbba39b1a5bd65b1bed3adec789d9e5168c Mon Sep 17 00:00:00 2001 From: jan Date: Thu, 3 Sep 2026 12:03:13 +0200 Subject: [PATCH 16/29] Refactor docstring in postprocess.py to simplify and clarify the purpose of the module --- learning_loop_node/detector/postprocess.py | 16 ++-------------- 1 file changed, 2 insertions(+), 14 deletions(-) diff --git a/learning_loop_node/detector/postprocess.py b/learning_loop_node/detector/postprocess.py index e544197b..fbfba50b 100644 --- a/learning_loop_node/detector/postprocess.py +++ b/learning_loop_node/detector/postprocess.py @@ -1,14 +1,4 @@ -"""Model-agnostic detection postprocessing. - -Every detection node ends up doing the same three things with a model's raw output: drop -low-confidence predictions, suppress overlapping boxes, and turn what survives into the -loop's detection dataclasses. None of that depends on the model. - -Two containers carry the same detections in this library: a detector node reports -:class:`~learning_loop_node.data_classes.image_metadata.ImageMetadata`, while a trainer's -auto-detection pass reports :class:`~learning_loop_node.data_classes.detections.Detections`. -:func:`to_image_metadata` and :func:`to_detections` build them from the same routine. -""" +"""Model-agnostic detection postprocessing.""" import logging from typing import NamedTuple @@ -29,7 +19,6 @@ logger = logging.getLogger(__name__) MIN_BOX_SIZE: int = 2 -"""A box is dropped unless both of its sides exceed this.""" class Detection(NamedTuple): @@ -148,8 +137,7 @@ def detections_from_xyxy( ) -> list[Detection]: """Convert already-suppressed model output into :class:`Detection` values. - For nodes whose model (or a torch/ONNX op) has done the suppression already, so only the - coordinate conversion is left. Corners are rounded here; :func:`post_process` truncates. + Corners are rounded here; :func:`post_process` truncates. """ result = [] for label, box, score in zip(labels, boxes, scores, strict=True): From 75208ba93e03ea9eec96f932d7bad7b2d4e4f58a Mon Sep 17 00:00:00 2001 From: Kevin Heye Date: Fri, 4 Sep 2026 14:17:01 +0200 Subject: [PATCH 17/29] Drop the centre-anchored clipping helper It came from Yolov5TrainerLogic.clip_box, whose only caller was the trainer's label-parsing pass. That caller now converts to the top-left convention before handing predictions to to_detections, which clips with clip_box -- so the centre-based variant would ship as public API with no consumer, for a convention the library has decided not to use internally. Co-Authored-By: Claude Opus 5 --- learning_loop_node/detector/geometry.py | 28 +------------------ .../tests/unit/test_geometry.py | 11 +------- 2 files changed, 2 insertions(+), 37 deletions(-) diff --git a/learning_loop_node/detector/geometry.py b/learning_loop_node/detector/geometry.py index db81abdc..479362ff 100644 --- a/learning_loop_node/detector/geometry.py +++ b/learning_loop_node/detector/geometry.py @@ -1,8 +1,7 @@ """Box and point clipping shared by every detector node. The loop stores a box as its top-left corner plus a size, which is the form :func:`clip_box` -produces. :func:`clip_box_centered` is the centre-anchored variant, for model outputs in that -convention. +produces. Centre-anchored model output converts before clipping. """ @@ -35,31 +34,6 @@ def clip_box( return clipped_x1, clipped_y1, clipped_width, clipped_height -def clip_box_centered( - *, - x: float, - y: float, - width: float, - height: float, - img_width: int, - img_height: int, -) -> tuple[float, float, float, float]: - """Clip a centre-anchored box to the image bounds, keeping it centre-anchored. - - Clipping moves the centre, because only the part of the box inside the image survives. - - :param x: Horizontal centre of the box. - :param y: Vertical centre of the box. - :return: The clipped ``(x, y, width, height)``, still centre-anchored. - """ - left = max(0.0, x - 0.5 * width) - top = max(0.0, y - 0.5 * height) - right = min(float(img_width), x + 0.5 * width) - bottom = min(float(img_height), y + 0.5 * height) - - return 0.5 * (left + right), 0.5 * (top + bottom), right - left, bottom - top - - def clip_point(x: float, y: float, img_width: int, img_height: int) -> tuple[float, float]: """Clamp a point into the image bounds.""" x = min(max(0, x), img_width) diff --git a/learning_loop_node/tests/unit/test_geometry.py b/learning_loop_node/tests/unit/test_geometry.py index f056f79d..a2f3836b 100644 --- a/learning_loop_node/tests/unit/test_geometry.py +++ b/learning_loop_node/tests/unit/test_geometry.py @@ -1,4 +1,4 @@ -from ...detector.geometry import clip_box, clip_box_centered, clip_point +from ...detector.geometry import clip_box, clip_point def test_box_inside_the_image_is_unchanged(): @@ -21,15 +21,6 @@ def test_box_corners_are_rounded(): assert clip_box(x1=10.4, y1=10.6, width=20.0, height=20.0, img_width=100, img_height=100) == (10, 11, 20, 20) -def test_centered_box_keeps_its_centre_when_it_fits(): - assert clip_box_centered(x=50, y=50, width=20, height=20, img_width=100, img_height=100) == (50, 50, 20, 20) - - -def test_clipping_a_centered_box_moves_its_centre(): - # only the right half of the box is inside the image, so the centre moves right - assert clip_box_centered(x=0, y=50, width=20, height=20, img_width=100, img_height=100) == (5, 50, 10, 20) - - def test_point_is_clamped_into_the_image(): assert clip_point(50, 50, 100, 100) == (50, 50) assert clip_point(-10, 150, 100, 100) == (0, 100) From b8e92948c34c29c501f4d85ece4ca35aef4ddac0 Mon Sep 17 00:00:00 2001 From: Kevin Heye Date: Fri, 4 Sep 2026 17:01:50 +0200 Subject: [PATCH 18/29] Share the trainer's GPU memory budgeting --vram-limit-gb is how a deployment shares one card between processes, and every trainer that probes for a batch size needs the same two halves of it: the budget the probe measures against, and the cap that holds the process to it. dfine had both; a second node would have copied them. Capping an allocator is a torch operation with no NVML equivalent, so this module imports torch. The package does not declare it: only a trainer imports the module, and a trainer brings torch already, so the library keeps installing on machines that train nothing. The unit test installs a stand-in under the name torch, which covers the arithmetic and the guards -- whether the cap holds is torch's own business. Co-Authored-By: Claude Opus 5 --- AGENTS.md | 7 ++ learning_loop_node/tests/unit/test_cuda.py | 90 ++++++++++++++++++++++ learning_loop_node/trainer/cuda.py | 65 ++++++++++++++++ 3 files changed, 162 insertions(+) create mode 100644 learning_loop_node/tests/unit/test_cuda.py create mode 100644 learning_loop_node/trainer/cuda.py diff --git a/AGENTS.md b/AGENTS.md index 9b03dc44..9a148a85 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -62,6 +62,13 @@ that fits, around a `fits` predicate the trainer supplies — none of it imports brings its own way of running a step. `macro_f1` scores the confusion matrix `_get_new_best_training_state` returns. +`trainer/cuda.py` is the one exception to that framework independence: `usable_memory_bytes` +and `limit_cuda_memory` turn a `--vram-limit-gb` setting into the budget a probe measures +against and the cap that holds the process to it, and capping an allocator has no NVML +equivalent. It imports torch, the package does **not** declare it, and only a trainer imports +the module — so the library keeps working where nothing trains. Its unit test installs a +stand-in under the name `torch`; whether the cap holds can only be seen on a card. + `helpers/entrypoint.py` holds what every node's `main.py` repeats: `node_parser` builds a configargparse parser with `--host`/`--port`, `run_node` starts uvicorn. A setting is a flag *and* an environment variable from one declaration — `--conf-threshold` reads diff --git a/learning_loop_node/tests/unit/test_cuda.py b/learning_loop_node/tests/unit/test_cuda.py new file mode 100644 index 00000000..8ec3ecd2 --- /dev/null +++ b/learning_loop_node/tests/unit/test_cuda.py @@ -0,0 +1,90 @@ +"""Tests for the one library module that imports torch. + +There is no torch in the dev extra, and no GPU in CI, so a stand-in is installed under the name +``torch`` before the module is imported. That covers the budgeting arithmetic and the guards -- +everything this module decides. Whether the cap actually holds is torch's own business and can +only be observed on a card. +""" + +import importlib +import logging +import sys +import types +from collections.abc import Callable +from typing import Any + +import pytest + +MODULE = 'learning_loop_node.trainer.cuda' +GIB = 1024**3 + + +@pytest.fixture(name='load') +def load_fixture(monkeypatch: pytest.MonkeyPatch) -> Callable[..., tuple[Any, list[tuple[float, int]]]]: + """Import the module against a fake torch; returns it and the fractions it capped with.""" + def load(*, cuda_available: bool = True, total_gb: float = 8.0): + capped: list[tuple[float, int]] = [] + torch = types.ModuleType('torch') + torch.cuda = types.SimpleNamespace( # type: ignore[attr-defined] + is_available=lambda: cuda_available, + empty_cache=lambda: capped.append((-1.0, -1)), + get_device_properties=lambda device: types.SimpleNamespace(total_memory=int(total_gb * GIB)), + set_per_process_memory_fraction=lambda fraction, device: capped.append((fraction, device)), + ) + monkeypatch.setitem(sys.modules, 'torch', torch) + monkeypatch.delitem(sys.modules, MODULE, raising=False) + return importlib.import_module(MODULE), capped + + yield load + sys.modules.pop(MODULE, None) + + +def test_no_limit_means_the_whole_card(load): + cuda, _ = load(total_gb=8.0) + assert cuda.usable_memory_bytes(0) == 8 * GIB + assert cuda.usable_memory_bytes(-1) == 8 * GIB + + +def test_a_limit_below_the_card_is_the_budget(load): + cuda, _ = load(total_gb=8.0) + assert cuda.usable_memory_bytes(6) == 6 * GIB + + +def test_a_limit_above_the_card_is_clamped_to_it(load): + # otherwise the safety margin would be a share of memory that does not exist + cuda, _ = load(total_gb=8.0) + assert cuda.usable_memory_bytes(16) == 8 * GIB + + +def test_the_cap_is_the_limits_share_of_the_card(load): + cuda, capped = load(total_gb=8.0) + cuda.limit_cuda_memory(2) + assert capped == [(0.25, 0)] + + +def test_no_limit_caps_nothing(load): + cuda, capped = load(total_gb=8.0) + cuda.limit_cuda_memory(0) + cuda.limit_cuda_memory(-1) + assert capped == [] + + +def test_nothing_is_capped_without_cuda(load): + cuda, capped = load(cuda_available=False) + cuda.limit_cuda_memory(2) + assert capped == [] + + +def test_a_limit_the_card_cannot_reach_warns_instead_of_capping(load, caplog): + # capping at a fraction >= 1 would be a no-op that reads as a limit having been applied + cuda, capped = load(total_gb=8.0) + with caplog.at_level(logging.WARNING): + cuda.limit_cuda_memory(8) + assert capped == [] + assert 'exceeds the card capacity' in caplog.text + + +def test_freeing_empties_the_cache(load): + cuda, capped = load() + cuda.free_cuda_memory() + assert capped == [(-1.0, -1)] diff --git a/learning_loop_node/trainer/cuda.py b/learning_loop_node/trainer/cuda.py new file mode 100644 index 00000000..25039e3b --- /dev/null +++ b/learning_loop_node/trainer/cuda.py @@ -0,0 +1,65 @@ +"""GPU memory budgeting for trainers that train in-process. + +A trainer that probes for its batch size needs two things a deployment can set: the budget the +probe measures against, and a cap that holds the process to it. Both come from one number -- +gigabytes of the card this training may use -- which is what a node exposes as +``--vram-limit-gb``. It is how one GPU gets shared between processes, and how a training keeps +headroom against allocator fragmentation. + +Unlike the rest of the library this module imports torch, and the package deliberately does not +declare it. Capping an allocator is a torch operation with no NVML equivalent, and declaring the +dependency would put an ML runtime into a library that is also installed on machines that train +nothing. Only a trainer imports this module, and a trainer brings torch already. + +Note the cap is relative to the card's *total* memory, not to what is free, so it does not +protect against another process claiming memory first. +""" +from __future__ import annotations + +import gc +import logging + +import torch + +logger = logging.getLogger(__name__) + + +def usable_memory_bytes(vram_limit_gb: float, device: int = 0) -> int: + """How much GPU memory this process may allocate, honouring :func:`limit_cuda_memory`. + + :param vram_limit_gb: 0 or less means the whole card. + """ + total_bytes = torch.cuda.get_device_properties(device).total_memory + if vram_limit_gb <= 0: + return total_bytes + return min(total_bytes, int(vram_limit_gb * 1024**3)) + + +def limit_cuda_memory(vram_limit_gb: float, device: int = 0) -> None: + """Cap how much of the GPU this process may allocate, to ``vram_limit_gb`` gigabytes. + + Call this once per process that touches the GPU -- a spawned training process included, + since the cap does not survive the spawn. + + :param vram_limit_gb: 0 or less means no cap. + """ + if vram_limit_gb <= 0 or not torch.cuda.is_available(): + return + + total_bytes = torch.cuda.get_device_properties(device).total_memory + fraction = vram_limit_gb * 1024**3 / total_bytes + total_gb = total_bytes / 1024**3 + + if fraction >= 1.0: + logger.warning('VRAM limit of %.1f GB exceeds the card capacity of %.1f GB; not limiting', + vram_limit_gb, total_gb) + return + + torch.cuda.set_per_process_memory_fraction(fraction, device) + logger.info('Limiting VRAM usage to %.1f GB of %.1f GB (%.0f%%)', vram_limit_gb, total_gb, fraction * 100) + + +def free_cuda_memory() -> None: + """Release cached CUDA memory, so a following allocation sees the real free space.""" + gc.collect() + torch.cuda.empty_cache() From f2f4c82b2856d917fa75b774b2ba0cc8d4c89cb2 Mon Sep 17 00:00:00 2001 From: Kevin Heye Date: Mon, 7 Sep 2026 08:46:59 +0200 Subject: [PATCH 19/29] Remove the pointer to a guide that does not exist The README linked docs/writing-a-node.md, which this branch never added. The guide is worth writing, but not as a passenger in a refactoring change. Co-Authored-By: Claude Opus 5 --- README.md | 2 -- 1 file changed, 2 deletions(-) diff --git a/README.md b/README.md index 0297b05f..19a58b84 100644 --- a/README.md +++ b/README.md @@ -13,8 +13,6 @@ This Python library helps to write Nodes that interact with the Zauberzeug Learn To start a node you have to implement the logic by inheriting from the corresponding base logic class. We provide samples in the 'mock' folders and recommend to follow that scheme. A complete trainer and detector example can be found [here](https://github.com/zauberzeug/yolov5_node). -**[docs/writing-a-node.md](docs/writing-a-node.md) walks through building a node from an empty repository**: which base class to implement per node type, what the library does for you, the trainer state machine, the trainer-to-detector model contract, and the repository and docker conventions the existing nodes share. The rest of this README documents how to *operate* a node. - #### Environment variables You can configure connection to our Learning Loop by specifying the following environment variables before starting: From 944c64a2204bf0f0ea2d542f9a294f775c7f46bf Mon Sep 17 00:00:00 2001 From: Kevin Heye Date: Mon, 7 Sep 2026 08:46:59 +0200 Subject: [PATCH 20/29] Spawn the iterator process explicitly, on every platform Without a context multiprocessing takes the platform default, which is fork on Linux -- and forking a process that has already initialised CUDA is what this helper exists to avoid. dfine enforces spawn globally today; now the helper does it wherever it runs, and picklability of the callable is a stated requirement rather than an accident of the platform. Co-Authored-By: Claude Opus 5 --- learning_loop_node/trainer/subprocess.py | 10 ++++++---- 1 file changed, 6 insertions(+), 4 deletions(-) diff --git a/learning_loop_node/trainer/subprocess.py b/learning_loop_node/trainer/subprocess.py index 4f038547..05386064 100644 --- a/learning_loop_node/trainer/subprocess.py +++ b/learning_loop_node/trainer/subprocess.py @@ -1,8 +1,9 @@ """Run a blocking, CPU-bound generator in its own process without blocking the event loop. -:func:`iterator_cpu_bound` runs the generator in a separate process and yields what it produces +:func:`iterator_cpu_bound` runs the generator in a spawned process and yields what it produces through a ``maxsize=1`` queue, so the producer never runs more than one item ahead of the -consumer. +consumer. The context is spawn on every platform, not the local default, so a process that has +already initialised CUDA is never forked; ``it`` and its arguments must therefore be picklable. Exceptions raised inside the process are re-raised in the caller, and the process is killed if the caller leaves the context early. @@ -42,8 +43,9 @@ async def _iterator_cpu_bound_inner( *args: P.args, **kwargs: P.kwargs, ) -> AsyncGenerator[T, None]: - state_queue: MPQueue[T | Exception | IteratorDone] = multiprocessing.Queue(maxsize=1) - process = multiprocessing.Process( + ctx = multiprocessing.get_context('spawn') + state_queue: MPQueue[T | Exception | IteratorDone] = ctx.Queue(maxsize=1) + process = ctx.Process( target=_iterator_wrapper, args=(it, state_queue, args, kwargs), name='iterator_cpu_bound', From 4ae56a59991b2cda962bb6661e955c28a6961541 Mon Sep 17 00:00:00 2001 From: Kevin Heye Date: Mon, 7 Sep 2026 08:46:59 +0200 Subject: [PATCH 21/29] Resolve uploaded categories through the shared lookup detector_node carried its own find_category_id_by_name, the second name lookup in the detector. category_by_name raises where this one returned '', so the caller catches: this path runs on metadata a client uploaded over socket.io, where an unknown name is the client's error and should cost that detection its id rather than the whole upload. The warning now names every unknown category once per upload instead of once per detection. Co-Authored-By: Claude Opus 5 --- learning_loop_node/detector/detector_node.py | 45 +++++++++++--------- 1 file changed, 24 insertions(+), 21 deletions(-) diff --git a/learning_loop_node/detector/detector_node.py b/learning_loop_node/detector/detector_node.py index eea3647d..3417069e 100644 --- a/learning_loop_node/detector/detector_node.py +++ b/learning_loop_node/detector/detector_node.py @@ -18,7 +18,6 @@ from ..data_classes import ( AboutResponse, - Category, Context, DetectorStatus, ImageMetadata, @@ -33,6 +32,7 @@ from ..helpers import background_tasks, environment_reader, run from ..helpers.misc import numpy_image_from_dict from ..node import Node +from .categories import category_by_name from .detector_logic import DetectorLogic, DetectorLogicFactory from .exceptions import NodeNeedsRestartError from .inbox_filter.relevance_filter import RelevanceFilter @@ -707,26 +707,29 @@ async def upload_images( await self.outbox.save(image, image_metadata, upload_priority) def add_category_id_to_detections(self, model_info: ModelInformation, image_metadata: ImageMetadata): - def find_category_id_by_name(categories: List[Category], category_name: str): - category_id = [category.id for category in categories if category.name == category_name] - return category_id[0] if category_id else '' - - for box_detection in image_metadata.box_detections: - category_name = box_detection.category_name - category_id = find_category_id_by_name(model_info.categories, category_name) - box_detection.category_id = category_id - for point_detection in image_metadata.point_detections: - category_name = point_detection.category_name - category_id = find_category_id_by_name(model_info.categories, category_name) - point_detection.category_id = category_id - for segmentation_detection in image_metadata.segmentation_detections: - category_name = segmentation_detection.category_name - category_id = find_category_id_by_name(model_info.categories, category_name) - segmentation_detection.category_id = category_id - for classification_detection in image_metadata.classification_detections: - category_name = classification_detection.category_name - category_id = find_category_id_by_name(model_info.categories, category_name) - classification_detection.category_id = category_id + """Resolve each detection's category id from its name, in metadata a client uploaded. + + A name the model does not know is the client's error, not a broken model, so it costs + that one detection its id rather than the whole upload. + """ + unknown_names: set[str] = set() + + def category_id_by_name(category_name: str) -> str: + try: + return category_by_name(model_info, category_name).id + except ValueError: + unknown_names.add(category_name) + return '' + + for detection in (*image_metadata.box_detections, + *image_metadata.point_detections, + *image_metadata.segmentation_detections, + *image_metadata.classification_detections): + detection.category_id = category_id_by_name(detection.category_name) + + if unknown_names: + self.log.warning('Model %s knows no category named %s', model_info.version, + ', '.join(sorted(unknown_names))) return image_metadata def register_sio_events(self, sio_client: AsyncClient): From f167834bdf956842afe706350c221a7b431b669f Mon Sep 17 00:00:00 2001 From: Kevin Heye Date: Mon, 7 Sep 2026 08:46:59 +0200 Subject: [PATCH 22/29] Test the unit suite on every supported Python version requires-python promises >=3.10, but 3.11 and 3.12 were tested nowhere -- and 3.12 is what every node but the Jetson detector runs. The suite needs no loop and no secrets, so covering all four costs only wall-clock on the job that gates the rest. --locked replaces --frozen so a lockfile that no longer matches pyproject.toml fails here rather than later, --python names the interpreter instead of taking whatever PATH offers first, and --no-sync keeps uv run from re-resolving into something other than what uv sync installed. Co-Authored-By: Claude Opus 5 --- .github/workflows/pytest.yml | 20 ++++++++++---------- 1 file changed, 10 insertions(+), 10 deletions(-) diff --git a/.github/workflows/pytest.yml b/.github/workflows/pytest.yml index e334d0b4..e6790960 100644 --- a/.github/workflows/pytest.yml +++ b/.github/workflows/pytest.yml @@ -6,21 +6,21 @@ jobs: unit: runs-on: ubuntu-latest timeout-minutes: 10 + strategy: + fail-fast: false + matrix: + python-version: ["3.10", "3.11", "3.12", "3.13"] steps: - uses: actions/checkout@v4 - - name: set up Python - uses: actions/setup-python@v5 - with: - python-version: "3.10" - - name: set up uv - uses: astral-sh/setup-uv@v3 + - uses: astral-sh/setup-uv@v3 - name: install dependencies - run: | - uv sync --extra dev --frozen + # --locked fails on a lockfile that no longer matches pyproject.toml + run: uv sync --extra dev --locked --python ${{ matrix.python-version }} + - name: check interpreter + run: uv run --no-sync python -c "import sys; assert sys.version_info[:2] == tuple(map(int, '${{ matrix.python-version }}'.split('.'))), sys.version" - name: test_unit # no Learning Loop and no secrets needed, so this suite gates the slow ones - run: | - uv run python -m pytest learning_loop_node/tests/unit -v + run: uv run --no-sync python -m pytest learning_loop_node/tests/unit -v pytest_3_10: needs: From f40ee9f77b7cc19f16676c00d1067f5a1808582b Mon Sep 17 00:00:00 2001 From: Kevin Heye Date: Mon, 7 Sep 2026 08:51:29 +0200 Subject: [PATCH 23/29] Name the model's output Prediction, and round its corners once Detection read like the base of BoxDetection, PointDetection and ClassificationDetection, and like the singular of Detections; it is neither. It is one surviving prediction in the model's own vocabulary -- a class index plus box geometry -- before category.type decides what it becomes. Prediction says that, category_index says which of the two identities it carries, and width/height/ confidence match the names the loop's dataclasses already use. It becomes a frozen keyword-only dataclass: four consecutive numbers built positionally make a swapped width/height type-correct and silent, and the detector hot path in yolov5 built it positionally. post_process now ends in predictions_from_xyxy instead of repeating the conversion, so there is one pipeline rather than two that had already drifted. Its int() went with the duplication: truncating here and rounding again in clip_box is strictly worse than rounding once, and the geometry stays float until the BoxDetection is built. Point detections gain real precision, since their centre no longer comes from truncated ints. Confidence is no longer rounded to two decimals on the way out -- that was never a contract, and the model's own score is what the loop should see. Co-Authored-By: Claude Opus 5 --- learning_loop_node/detector/postprocess.py | 110 ++++++++---------- .../tests/unit/test_postprocess.py | 39 ++++--- 2 files changed, 73 insertions(+), 76 deletions(-) diff --git a/learning_loop_node/detector/postprocess.py b/learning_loop_node/detector/postprocess.py index fbfba50b..a6812ddc 100644 --- a/learning_loop_node/detector/postprocess.py +++ b/learning_loop_node/detector/postprocess.py @@ -1,7 +1,8 @@ """Model-agnostic detection postprocessing.""" import logging -from typing import NamedTuple +from collections.abc import Sequence +from dataclasses import dataclass import numpy as np @@ -21,16 +22,17 @@ MIN_BOX_SIZE: int = 2 -class Detection(NamedTuple): - """One surviving prediction. ``x``/``y`` are the top-left corner, ``category`` is an index - into :attr:`ModelInformation.categories`.""" +@dataclass(kw_only=True, slots=True, frozen=True) +class Prediction: + """One surviving prediction in the model's coordinates: top-left corner and size in pixels, + unrounded, and the category as an index into :attr:`ModelInformation.categories`.""" - x: int - y: int - w: int - h: int - category: int - probability: float + x: float + y: float + width: float + height: float + category_index: int + confidence: float def post_process( @@ -42,8 +44,8 @@ def post_process( iou_threshold: float, origin_h: int, origin_w: int, -) -> list[Detection]: - """Filter by confidence, run NMS, return a :class:`Detection` list in x/y/w/h form.""" +) -> list[Prediction]: + """Filter by confidence, run NMS, return what survives.""" mask = scores > conf_threshold boxes = boxes[mask].copy() scores = scores[mask] @@ -56,14 +58,7 @@ def post_process( boxes, scores, classes, iou_threshold=iou_threshold, origin_h=origin_h, origin_w=origin_w) - result = [] - for j, box in enumerate(boxes): - x1, y1, x2, y2 = box - w = x2 - x1 - h = y2 - y1 - result.append(Detection(x=int(x1), y=int(y1), w=int(w), h=int(h), - category=int(classes[j]), probability=round(float(scores[j]), 2))) - return result + return predictions_from_xyxy(labels=classes, boxes=boxes, scores=scores) def non_max_suppression( @@ -129,38 +124,36 @@ def bbox_iou( return inter_area / (b1_area + b2_area - inter_area + 1e-16) -def detections_from_xyxy( +def predictions_from_xyxy( *, - labels: list[float], - boxes: list[list[float]], - scores: list[float], -) -> list[Detection]: - """Convert already-suppressed model output into :class:`Detection` values. - - Corners are rounded here; :func:`post_process` truncates. + labels: Sequence[float], + boxes: Sequence[Sequence[float]], + scores: Sequence[float], +) -> list[Prediction]: + """Convert xyxy model output into predictions, for models that suppress their own overlaps. + + Corners stay unrounded: :func:`clip_box` rounds once, when the box becomes a + :class:`BoxDetection`. """ - result = [] - for label, box, score in zip(labels, boxes, scores, strict=True): - x1, y1, x2, y2 = (round(value) for value in box) - result.append(Detection(x=x1, y=y1, w=x2 - x1, h=y2 - y1, - category=int(label), probability=score)) - return result + return [Prediction(x=float(x1), y=float(y1), width=float(x2 - x1), height=float(y2 - y1), + category_index=int(label), confidence=float(score)) + for label, (x1, y1, x2, y2), score in zip(labels, boxes, scores, strict=True)] def to_image_metadata( - detections: list[Detection], + predictions: list[Prediction], model_information: ModelInformation, im_height: int, im_width: int, ) -> ImageMetadata: - """Build the container a *detector* node reports from a list of detections.""" + """Build the container a *detector* node reports.""" image_metadata = ImageMetadata() - _append_detections(image_metadata, detections, model_information, im_height, im_width) + _append_predictions(image_metadata, predictions, model_information, im_height, im_width) return image_metadata def to_detections( - detections: list[Detection], + predictions: list[Prediction], model_information: ModelInformation, im_height: int, im_width: int, @@ -169,32 +162,31 @@ def to_detections( ) -> Detections: """Build the container a *trainer*'s auto-detection pass reports.""" result = Detections(image_id=image_id) - _append_detections(result, detections, model_information, im_height, im_width) + _append_predictions(result, predictions, model_information, im_height, im_width) return result -def _append_detections( +def _append_predictions( target: ImageMetadata | Detections, - detections: list[Detection], + predictions: list[Prediction], model_information: ModelInformation, im_height: int, im_width: int, ) -> None: - """Resolve each detection's category and append it to ``target``, clipped to the image.""" - skipped_detections = [] - - for detection in detections: - x, y, w, h, category_idx, probability = detection - category = category_by_index(model_information, category_idx) - if w <= MIN_BOX_SIZE or h <= MIN_BOX_SIZE: - skipped_detections.append((category.name, detection)) + """Resolve each prediction's category and append it to ``target``, clipped to the image.""" + skipped_predictions = [] + + for prediction in predictions: + category = category_by_index(model_information, prediction.category_index) + if prediction.width <= MIN_BOX_SIZE or prediction.height <= MIN_BOX_SIZE: + skipped_predictions.append((category.name, prediction)) continue if category.type == CategoryType.Box: clipped_x1, clipped_y1, clipped_w, clipped_h = clip_box( - x1=x, - y1=y, - width=w, - height=h, + x1=prediction.x, + y1=prediction.y, + width=prediction.width, + height=prediction.height, img_width=im_width, img_height=im_height, ) @@ -207,11 +199,11 @@ def _append_detections( height=clipped_h, category_id=category.id, model_name=model_information.version, - confidence=probability, + confidence=prediction.confidence, ) ) elif category.type == CategoryType.Point: - cx, cy = x + w / 2, y + h / 2 + cx, cy = prediction.x + prediction.width / 2, prediction.y + prediction.height / 2 cx, cy = clip_point(cx, cy, im_width, im_height) target.point_detections.append( PointDetection( @@ -220,12 +212,12 @@ def _append_detections( y=cy, category_id=category.id, model_name=model_information.version, - confidence=probability, + confidence=prediction.confidence, ) ) else: logger.warning('Unsupported category type %s for category %s', category.type, category.name) - if skipped_detections: - log_msg = '\n'.join([str(d) for d in skipped_detections]) - logger.warning('Removed %d small detections from result: \n%s', len(skipped_detections), log_msg) + if skipped_predictions: + log_msg = '\n'.join([str(p) for p in skipped_predictions]) + logger.warning('Removed %d small detections from result: \n%s', len(skipped_predictions), log_msg) diff --git a/learning_loop_node/tests/unit/test_postprocess.py b/learning_loop_node/tests/unit/test_postprocess.py index c70562a0..6d3ddcf4 100644 --- a/learning_loop_node/tests/unit/test_postprocess.py +++ b/learning_loop_node/tests/unit/test_postprocess.py @@ -3,11 +3,11 @@ from ...data_classes import Category, ModelInformation from ...detector.postprocess import ( - Detection, + Prediction, bbox_iou, - detections_from_xyxy, non_max_suppression, post_process, + predictions_from_xyxy, to_detections, to_image_metadata, ) @@ -66,7 +66,12 @@ def test_post_process_drops_predictions_below_the_confidence_threshold(): boxes = np.array([[10, 10, 60, 60], [100, 100, 150, 150]], dtype=np.float32) result = post_process(boxes, np.array([0.9, 0.1], dtype=np.float32), np.array([0, 0]), conf_threshold=0.5, iou_threshold=0.45, origin_h=200, origin_w=200) - assert result == [Detection(10, 10, 50, 50, 0, 0.9)] + assert len(result) == 1 + prediction = result[0] + assert (prediction.x, prediction.y, prediction.width, prediction.height) == (10, 10, 50, 50) + assert prediction.category_index == 0 + # the model's own float32 score, no longer rounded to two decimals on the way out + assert prediction.confidence == pytest.approx(0.9) def test_post_process_on_an_empty_prediction_returns_nothing(): @@ -75,20 +80,20 @@ def test_post_process_on_an_empty_prediction_returns_nothing(): conf_threshold=0.5, iou_threshold=0.45, origin_h=10, origin_w=10) == [] -def test_already_suppressed_output_is_converted_with_rounded_corners(): - assert detections_from_xyxy(labels=[1.0], boxes=[[10.4, 10.6, 60.4, 60.6]], scores=[0.55]) == \ - [Detection(10, 11, 50, 50, 1, 0.55)] +def test_already_suppressed_output_keeps_the_model_s_own_precision(): + assert predictions_from_xyxy(labels=[1.0], boxes=[[10.4, 10.6, 60.4, 60.6]], scores=[0.55]) == \ + [Prediction(x=10.4, y=10.6, width=50.0, height=50.0, category_index=1, confidence=0.55)] def test_converting_already_suppressed_output_requires_matching_lengths(): with pytest.raises(ValueError): - detections_from_xyxy(labels=[1.0, 2.0], boxes=[[0.0, 0.0, 1.0, 1.0]], scores=[0.5]) + predictions_from_xyxy(labels=[1.0, 2.0], boxes=[[0.0, 0.0, 1.0, 1.0]], scores=[0.5]) # ---------------------------------------------------------------- building the containers def test_a_box_category_becomes_a_box_detection(): - metadata = to_image_metadata([Detection(10, 20, 30, 40, 0, 0.9)], model_information(), 200, 200) + metadata = to_image_metadata([Prediction(x=10, y=20, width=30, height=40, category_index=0, confidence=0.9)], model_information(), 200, 200) assert len(metadata.point_detections) == 0 detection = metadata.box_detections[0] assert (detection.x, detection.y, detection.width, detection.height) == (10, 20, 30, 40) @@ -98,7 +103,7 @@ def test_a_box_category_becomes_a_box_detection(): def test_a_point_category_becomes_the_centre_of_the_box(): - metadata = to_image_metadata([Detection(100, 100, 40, 40, 1, 0.7)], model_information(), 200, 200) + metadata = to_image_metadata([Prediction(x=100, y=100, width=40, height=40, category_index=1, confidence=0.7)], model_information(), 200, 200) assert len(metadata.box_detections) == 0 detection = metadata.point_detections[0] assert (detection.x, detection.y) == (120, 120) @@ -106,26 +111,26 @@ def test_a_point_category_becomes_the_centre_of_the_box(): def test_detections_are_clipped_to_the_image(): - metadata = to_image_metadata([Detection(-20, -20, 60, 60, 0, 0.5)], model_information(), 200, 200) + metadata = to_image_metadata([Prediction(x=-20, y=-20, width=60, height=60, category_index=0, confidence=0.5)], model_information(), 200, 200) detection = metadata.box_detections[0] assert (detection.x, detection.y, detection.width, detection.height) == (0, 0, 40, 40) @pytest.mark.parametrize('width,height', [(2, 30), (30, 2), (1, 1)]) def test_boxes_too_small_to_be_useful_are_dropped(width: int, height: int): - metadata = to_image_metadata([Detection(5, 5, width, height, 0, 0.5)], model_information(), 200, 200) + metadata = to_image_metadata([Prediction(x=5, y=5, width=width, height=height, category_index=0, confidence=0.5)], model_information(), 200, 200) assert len(metadata) == 0 def test_a_category_type_the_node_cannot_report_is_skipped(): classification = Category(id='uuid-cls', name='ripe', type=CategoryType.Classification) - metadata = to_image_metadata([Detection(10, 10, 30, 30, 0, 0.5)], + metadata = to_image_metadata([Prediction(x=10, y=10, width=30, height=30, category_index=0, confidence=0.5)], model_information(classification), 200, 200) assert len(metadata) == 0 def test_the_trainer_container_carries_the_image_id(): - result = to_detections([Detection(10, 20, 30, 40, 0, 0.9)], model_information(), 200, 200, + result = to_detections([Prediction(x=10, y=20, width=30, height=40, category_index=0, confidence=0.9)], model_information(), 200, 200, image_id='image-uuid') assert result.image_id == 'image-uuid' assert len(result.box_detections) == 1 @@ -133,10 +138,10 @@ def test_the_trainer_container_carries_the_image_id(): def test_trainer_and_detector_paths_agree_on_the_same_detections(): """The whole point of sharing this code: auto-detections and live detections must match.""" - detections = [Detection(-5, -5, 60, 60, 0, 0.9), Detection(100, 100, 40, 40, 1, 0.7), - Detection(5, 5, 1, 1, 0, 0.5)] - metadata = to_image_metadata(detections, model_information(), 200, 200) - result = to_detections(detections, model_information(), 200, 200, image_id='image-uuid') + predictions = [Prediction(x=-5, y=-5, width=60, height=60, category_index=0, confidence=0.9), Prediction(x=100, y=100, width=40, height=40, category_index=1, confidence=0.7), + Prediction(x=5, y=5, width=1, height=1, category_index=0, confidence=0.5)] + metadata = to_image_metadata(predictions, model_information(), 200, 200) + result = to_detections(predictions, model_information(), 200, 200, image_id='image-uuid') assert [(d.x, d.y, d.width, d.height, d.category_id) for d in result.box_detections] == \ [(d.x, d.y, d.width, d.height, d.category_id) for d in metadata.box_detections] From 74b914f532b0a2553764bcf7d092b6421b6b8960 Mon Sep 17 00:00:00 2001 From: Kevin Heye Date: Mon, 7 Sep 2026 09:05:41 +0200 Subject: [PATCH 24/29] Run every suite on both ends of the supported range The two jobs split the suites between versions rather than the versions between suites: general, detector and mock_detector only ever ran on 3.10, annotator, trainer and mock_trainer only on 3.13. Neither version was actually tested, only half of each. One matrix job runs all six on 3.10 and 3.13 -- the ends of what requires-python promises, with the unit job covering 3.11 and 3.12 as well. max-parallel keeps the entries serialised: they share one Learning Loop, and test_general creates and deletes zauberzeug/pytest_nodelib_general by a fixed name, so two entries at once would delete each other's project. That serialisation used to come from chaining the two jobs with needs. Every suite collects on both versions locally; whether they pass needs the loop. Co-Authored-By: Claude Opus 5 --- .github/workflows/pytest.yml | 66 +++++++++++------------------------- 1 file changed, 19 insertions(+), 47 deletions(-) diff --git a/.github/workflows/pytest.yml b/.github/workflows/pytest.yml index e6790960..929eb642 100644 --- a/.github/workflows/pytest.yml +++ b/.github/workflows/pytest.yml @@ -22,99 +22,71 @@ jobs: # no Learning Loop and no secrets needed, so this suite gates the slow ones run: uv run --no-sync python -m pytest learning_loop_node/tests/unit -v - pytest_3_10: + pytest: needs: - unit runs-on: ubuntu-latest - timeout-minutes: 30 + timeout-minutes: 45 strategy: fail-fast: false + # these suites share one Learning Loop instance, and test_general creates and deletes + # the project zauberzeug/pytest_nodelib_general by a fixed name -- two entries at once + # would delete each other's project + max-parallel: 1 + matrix: + # the ends of what requires-python promises; the unit job covers 3.11 and 3.12 too, + # and no node runs anything outside this range + python-version: ["3.10", "3.13"] steps: - uses: actions/checkout@v4 - - name: set up Python - uses: actions/setup-python@v5 - with: - python-version: "3.10" - - name: set up uv - uses: astral-sh/setup-uv@v3 + - uses: astral-sh/setup-uv@v3 - name: install test dependencies run: | sudo apt-get update sudo apt-get install libcurl4-openssl-dev libssl-dev jpeginfo - name: install dependencies - run: | - uv sync --extra dev --frozen + run: uv sync --extra dev --locked --python ${{ matrix.python-version }} - name: test_general env: LOOP_HOST: "preview.learning-loop.ai" LOOP_USERNAME: "admin" LOOP_PASSWORD: ${{ secrets.LEARNING_LOOP_ADMIN_PASSWORD }} - run: | - uv run python -m pytest "learning_loop_node/tests/general" -v + run: uv run --no-sync python -m pytest learning_loop_node/tests/general -v - name: test_detector env: LOOP_HOST: "preview.learning-loop.ai" LOOP_USERNAME: "admin" LOOP_PASSWORD: ${{ secrets.LEARNING_LOOP_ADMIN_PASSWORD }} - run: | - uv run python -m pytest learning_loop_node/tests/detector -v + run: uv run --no-sync python -m pytest learning_loop_node/tests/detector -v - name: test_mock_detector env: LOOP_HOST: "preview.learning-loop.ai" LOOP_USERNAME: "admin" LOOP_PASSWORD: ${{ secrets.LEARNING_LOOP_ADMIN_PASSWORD }} - run: | - uv run python -m pytest mock_detector -v - - pytest_3_13: - needs: - - pytest_3_10 - runs-on: ubuntu-latest - timeout-minutes: 30 - strategy: - fail-fast: false - steps: - - uses: actions/checkout@v4 - - name: set up Python - uses: actions/setup-python@v5 - with: - python-version: "3.13" - - name: set up uv - uses: astral-sh/setup-uv@v3 - - name: install test dependencies - run: | - sudo apt-get update - sudo apt-get install libcurl4-openssl-dev libssl-dev jpeginfo - - name: install dependencies - run: | - uv sync --extra dev --frozen + run: uv run --no-sync python -m pytest mock_detector -v - name: test_annotator env: LOOP_HOST: "preview.learning-loop.ai" LOOP_USERNAME: "admin" LOOP_PASSWORD: ${{ secrets.LEARNING_LOOP_ADMIN_PASSWORD }} - run: | - uv run python -m pytest learning_loop_node/tests/annotator -v + run: uv run --no-sync python -m pytest learning_loop_node/tests/annotator -v - name: test_trainer env: LOOP_HOST: "preview.learning-loop.ai" LOOP_USERNAME: "admin" LOOP_PASSWORD: ${{ secrets.LEARNING_LOOP_ADMIN_PASSWORD }} - run: | - uv run python -m pytest learning_loop_node/tests/trainer -v + run: uv run --no-sync python -m pytest learning_loop_node/tests/trainer -v - name: test_mock_trainer env: LOOP_HOST: "preview.learning-loop.ai" LOOP_USERNAME: "admin" LOOP_PASSWORD: ${{ secrets.LEARNING_LOOP_ADMIN_PASSWORD }} - run: | - uv run python -m pytest mock_trainer -v + run: uv run --no-sync python -m pytest mock_trainer -v slack: needs: - unit - - pytest_3_10 - - pytest_3_13 + - pytest if: always() # also execute when pytest fails runs-on: ubuntu-latest steps: From 4dfe4deb838f187314d1a5b13376f09ee51e7c0d Mon Sep 17 00:00:00 2001 From: Kevin Heye Date: Mon, 7 Sep 2026 12:59:15 +0200 Subject: [PATCH 25/29] Share the torch side of a batch-size probe, not just the search find_batch_size arrived here without the measuring around it, so a node adopting it still had to write the part that is easy to get wrong. dfine_node had written it twice over -- once for training, once for the detection pass -- and what those two share is not model-specific at all: hold back a safety margin, reset the peak counters, run the step, and decide whether a failure means "too big" or "broken". probe_batch_size is the whole probe for a node whose measurement is a single call. It resolves the limit == 0 sentinel, falls back without a card, holds the margin for the length of the search and logs the size it settled on. measured_fits is that skeleton on its own, for a probe that has to build a throwaway model before it can measure: the margin has to be claimed before the model is built, so such a probe composes reserve_margin, measured_fits and find_batch_size itself rather than handing over a closure. on_out_of_memory is what lets it drop an optimizer's gradients after a failed trial without also owning the guard. The guard is the reason this belongs in the library rather than in each node. isinstance(exc, torch.cuda.OutOfMemoryError) or is_out_of_memory(exc) is one line, and a node that writes `except RuntimeError: return False` instead reports the smallest batch size as the card's fault whenever a shape is wrong. Both repositories that had a probe had their own private copy of that predicate. MAX_BATCH_SIZE moves in beside NO_GPU_BATCH_SIZE and MIN_TRAIN_STEPS_PER_EPOCH: it is search policy, not a torch concern, and a node resolving its own sentinel wants it from the framework-free module. SAFETY_MARGIN stays here, since claiming the margin is an allocation. 14 tests against the same stand-in for torch, covering what the module decides: the doubling and the power-of-two rounding, the margin as a share of the budget rather than of the card, the fallback that runs nothing without a GPU, MemoryError and the bare RuntimeErrors that cuDNN and cuBLAS raise -- and that a failure which is not about memory propagates rather than reading as a full card. Co-Authored-By: Claude Opus 5 --- AGENTS.md | 21 ++- learning_loop_node/tests/unit/test_cuda.py | 209 ++++++++++++++++++--- learning_loop_node/trainer/batch_size.py | 3 + learning_loop_node/trainer/cuda.py | 96 +++++++++- 4 files changed, 298 insertions(+), 31 deletions(-) diff --git a/AGENTS.md b/AGENTS.md index 9a148a85..716853ff 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -62,12 +62,21 @@ that fits, around a `fits` predicate the trainer supplies — none of it imports brings its own way of running a step. `macro_f1` scores the confusion matrix `_get_new_best_training_state` returns. -`trainer/cuda.py` is the one exception to that framework independence: `usable_memory_bytes` -and `limit_cuda_memory` turn a `--vram-limit-gb` setting into the budget a probe measures -against and the cap that holds the process to it, and capping an allocator has no NVML -equivalent. It imports torch, the package does **not** declare it, and only a trainer imports -the module — so the library keeps working where nothing trains. Its unit test installs a -stand-in under the name `torch`; whether the cap holds can only be seen on a card. +`trainer/cuda.py` is the one exception to that framework independence, and holds everything +about a batch-size probe that torch has to answer. `usable_memory_bytes` and `limit_cuda_memory` +turn a `--vram-limit-gb` setting into the budget a probe measures against and the cap that holds +the process to it, and capping an allocator has no NVML equivalent. `probe_batch_size` is the +whole probe for a node whose measurement is a single call — it resolves the limit, falls back +without a card, holds the safety margin and runs the search. A node that must build a throwaway +model first reserves the margin before building it, and so composes the same pieces itself: +`reserve_margin`, `measured_fits` and `find_batch_size`. `measured_fits` is where an +out-of-memory failure is told from a bug — both arrive as the same exception types, and a probe +that confuses them reports the smallest batch size as the card's fault. + +It imports torch, the package does **not** declare it, and only a trainer imports the module — so +the library keeps working where nothing trains. Its unit test installs a stand-in under the name +`torch`, which covers the arithmetic, the guards and the search; whether the cap holds, and what +a real step costs, can only be seen on a card. `helpers/entrypoint.py` holds what every node's `main.py` repeats: `node_parser` builds a configargparse parser with `--host`/`--port`, `run_node` starts uvicorn. A setting is a flag diff --git a/learning_loop_node/tests/unit/test_cuda.py b/learning_loop_node/tests/unit/test_cuda.py index 8ec3ecd2..d7daad0a 100644 --- a/learning_loop_node/tests/unit/test_cuda.py +++ b/learning_loop_node/tests/unit/test_cuda.py @@ -1,9 +1,9 @@ """Tests for the one library module that imports torch. There is no torch in the dev extra, and no GPU in CI, so a stand-in is installed under the name -``torch`` before the module is imported. That covers the budgeting arithmetic and the guards -- -everything this module decides. Whether the cap actually holds is torch's own business and can -only be observed on a card. +``torch`` before the module is imported. That covers the budgeting arithmetic, the guards and the +search around a probe -- everything this module decides. Whether the cap actually holds, and what +a real step costs, is torch's own business and can only be observed on a card. """ import importlib @@ -15,30 +15,68 @@ import pytest +from learning_loop_node.trainer.batch_size import MAX_BATCH_SIZE, NO_GPU_BATCH_SIZE +from learning_loop_node.trainer.exceptions import InsufficientMemoryError + MODULE = 'learning_loop_node.trainer.cuda' GIB = 1024**3 -@pytest.fixture(name='load') -def load_fixture(monkeypatch: pytest.MonkeyPatch) -> Callable[..., tuple[Any, list[tuple[float, int]]]]: - """Import the module against a fake torch; returns it and the fractions it capped with.""" - def load(*, cuda_available: bool = True, total_gb: float = 8.0): - capped: list[tuple[float, int]] = [] - torch = types.ModuleType('torch') - torch.cuda = types.SimpleNamespace( # type: ignore[attr-defined] +class _FakeTorch: + """What the module uses of torch, plus a record of what it asked for.""" + + def __init__(self, *, cuda_available: bool, total_gb: float, peak_gb: float) -> None: + self.capped: list[tuple[float, int]] = [] + self.allocated: list[int] = [] + + class OutOfMemoryError(RuntimeError): + """Torch's own; note the module may not rely on its message.""" + + self.OutOfMemoryError = OutOfMemoryError # named as torch spells it + self.module = types.ModuleType('torch') + self.module.uint8 = 'uint8' # type: ignore[attr-defined] + self.module.empty = self._empty # type: ignore[attr-defined] + self.module.cuda = types.SimpleNamespace( # type: ignore[attr-defined] + OutOfMemoryError=OutOfMemoryError, is_available=lambda: cuda_available, - empty_cache=lambda: capped.append((-1.0, -1)), + empty_cache=lambda: self.capped.append((-1.0, -1)), get_device_properties=lambda device: types.SimpleNamespace(total_memory=int(total_gb * GIB)), - set_per_process_memory_fraction=lambda fraction, device: capped.append((fraction, device)), + set_per_process_memory_fraction=lambda fraction, device: self.capped.append((fraction, device)), + reset_peak_memory_stats=lambda: None, + max_memory_allocated=lambda: int(peak_gb * GIB), + synchronize=lambda: None, ) - monkeypatch.setitem(sys.modules, 'torch', torch) + + def _empty(self, count: int, dtype: str, device: str) -> object: + assert (dtype, device) == ('uint8', 'cuda') + self.allocated.append(count) + return object() + + +@pytest.fixture(name='load') +def load_fixture(monkeypatch: pytest.MonkeyPatch) -> Callable[..., tuple[Any, _FakeTorch]]: + """Import the module against a fake torch; returns it and the stand-in it ran against.""" + def load(*, cuda_available: bool = True, total_gb: float = 8.0, peak_gb: float = 2.0): + fake = _FakeTorch(cuda_available=cuda_available, total_gb=total_gb, peak_gb=peak_gb) + monkeypatch.setitem(sys.modules, 'torch', fake.module) monkeypatch.delitem(sys.modules, MODULE, raising=False) - return importlib.import_module(MODULE), capped + return importlib.import_module(MODULE), fake yield load sys.modules.pop(MODULE, None) +def _fits_up_to(largest: int, fake: _FakeTorch, ran: list[int]) -> Callable[[int], None]: + """A step that runs out of memory above ``largest``, recording every size it was asked for.""" + def run_batch(batch_size: int) -> None: + ran.append(batch_size) + if batch_size > largest: + raise fake.OutOfMemoryError('tried to allocate 20.00 GiB') + return run_batch + + +# --- the budget and the cap --- + def test_no_limit_means_the_whole_card(load): cuda, _ = load(total_gb=8.0) assert cuda.usable_memory_bytes(0) == 8 * GIB @@ -57,34 +95,157 @@ def test_a_limit_above_the_card_is_clamped_to_it(load): def test_the_cap_is_the_limits_share_of_the_card(load): - cuda, capped = load(total_gb=8.0) + cuda, fake = load(total_gb=8.0) cuda.limit_cuda_memory(2) - assert capped == [(0.25, 0)] + assert fake.capped == [(0.25, 0)] def test_no_limit_caps_nothing(load): - cuda, capped = load(total_gb=8.0) + cuda, fake = load(total_gb=8.0) cuda.limit_cuda_memory(0) cuda.limit_cuda_memory(-1) - assert capped == [] + assert fake.capped == [] def test_nothing_is_capped_without_cuda(load): - cuda, capped = load(cuda_available=False) + cuda, fake = load(cuda_available=False) cuda.limit_cuda_memory(2) - assert capped == [] + assert fake.capped == [] def test_a_limit_the_card_cannot_reach_warns_instead_of_capping(load, caplog): # capping at a fraction >= 1 would be a no-op that reads as a limit having been applied - cuda, capped = load(total_gb=8.0) + cuda, fake = load(total_gb=8.0) with caplog.at_level(logging.WARNING): cuda.limit_cuda_memory(8) - assert capped == [] + assert fake.capped == [] assert 'exceeds the card capacity' in caplog.text def test_freeing_empties_the_cache(load): - cuda, capped = load() + cuda, fake = load() cuda.free_cuda_memory() - assert capped == [(-1.0, -1)] + assert fake.capped == [(-1.0, -1)] + + +# --- the safety margin --- + +def test_the_margin_is_a_share_of_the_budget(load): + cuda, fake = load(total_gb=8.0) + cuda.reserve_margin(4, probe='probe') + assert fake.allocated == [int(4 * GIB * cuda.SAFETY_MARGIN)] + + +def test_the_margin_is_a_share_of_the_whole_card_when_nothing_is_budgeted(load): + cuda, fake = load(total_gb=8.0) + cuda.reserve_margin(0, probe='probe') + assert fake.allocated == [int(8 * GIB * cuda.SAFETY_MARGIN)] + + +# --- probing --- + +def test_the_probe_keeps_the_largest_size_that_fits(load): + cuda, fake = load() + ran: list[int] = [] + assert cuda.probe_batch_size(_fits_up_to(16, fake, ran), limit=64) == 16 + assert ran == [1, 2, 4, 8, 16, 32], 'only powers of two, and one trial past the answer' + + +def test_the_probe_reserves_the_margin_before_it_measures(load): + # a trial that ran against the full card would choose a size the training cannot keep + cuda, fake = load(total_gb=8.0) + + def run_batch(_: int) -> None: + assert fake.allocated == [int(8 * GIB * cuda.SAFETY_MARGIN)] + + cuda.probe_batch_size(run_batch, limit=2) + + +def test_the_limit_is_rounded_down_to_a_power_of_two(load): + cuda, fake = load() + ran: list[int] = [] + assert cuda.probe_batch_size(_fits_up_to(1024, fake, ran), limit=12) == 8 + + +def test_an_unset_limit_stops_at_the_maximum(load): + cuda, fake = load() + ran: list[int] = [] + assert cuda.probe_batch_size(_fits_up_to(2048, fake, ran)) == MAX_BATCH_SIZE + + +def test_a_probe_without_a_gpu_does_not_run_the_step(load): + """With no card to measure, the fallback caps the batch and never runs the step.""" + cuda, fake = load(cuda_available=False) + ran: list[int] = [] + assert cuda.probe_batch_size(_fits_up_to(1024, fake, ran), limit=32) == NO_GPU_BATCH_SIZE + assert cuda.probe_batch_size(_fits_up_to(1024, fake, ran), limit=2) == 2 + assert not ran, 'nothing may be run without a GPU' + assert not fake.allocated, 'and no margin claimed on a card that is not there' + + +def test_a_batch_size_of_one_that_does_not_fit_is_an_error(load): + cuda, fake = load() + ran: list[int] = [] + with pytest.raises(InsufficientMemoryError): + cuda.probe_batch_size(_fits_up_to(0, fake, ran), limit=32) + + +def test_exhausted_host_memory_also_means_it_does_not_fit(load): + cuda, _ = load() + + def run_batch(batch_size: int) -> None: + if batch_size > 2: + raise MemoryError + + assert cuda.probe_batch_size(run_batch, limit=32) == 2 + + +def test_a_plain_runtime_error_saying_it_is_out_of_memory_does_not_fit(load): + # cuDNN and cuBLAS workspaces arrive as bare RuntimeErrors + cuda, _ = load() + + def run_batch(batch_size: int) -> None: + if batch_size > 2: + raise RuntimeError('cuDNN error: CUDNN_STATUS_ALLOC_FAILED') + + assert cuda.probe_batch_size(run_batch, limit=32) == 2 + + +def test_torchs_own_error_needs_no_recognisable_message(load): + cuda, fake = load() + + def run_batch(batch_size: int) -> None: + if batch_size > 2: + raise fake.OutOfMemoryError('see the memory summary') + + assert cuda.probe_batch_size(run_batch, limit=32) == 2 + + +def test_a_failure_that_is_not_about_memory_is_a_bug_and_propagates(load): + # a probe that swallows this would report the smallest batch size as the card's fault + cuda, _ = load() + + def run_batch(batch_size: int) -> None: + if batch_size > 2: + raise RuntimeError('mat1 and mat2 shapes cannot be multiplied') + + with pytest.raises(RuntimeError, match='shapes cannot be multiplied'): + cuda.probe_batch_size(run_batch, limit=32) + + +def test_what_the_step_reports_is_logged_beside_the_peak(load, caplog): + cuda, _ = load(peak_gb=3.0) + with caplog.at_level(logging.INFO): + cuda.probe_batch_size(lambda _: 'validation peaked at 1.00 GB', probe='miniature epoch', limit=1) + assert 'miniature epoch: 1 fits (peak 3.00 GB, margin included); validation peaked at 1.00 GB' in caplog.text + + +def test_only_a_trial_that_ran_out_of_memory_gets_cleaned_up_after(load): + cuda, fake = load() + cleaned: list[int] = [] + ran: list[int] = [] + fits = cuda.measured_fits(_fits_up_to(2, fake, ran), probe='probe', + on_out_of_memory=lambda: cleaned.append(len(ran))) + + assert [fits(size) for size in (1, 2, 4)] == [True, True, False] + assert cleaned == [3], 'once, after the third trial' diff --git a/learning_loop_node/trainer/batch_size.py b/learning_loop_node/trainer/batch_size.py index 0e438036..78b46f9c 100644 --- a/learning_loop_node/trainer/batch_size.py +++ b/learning_loop_node/trainer/batch_size.py @@ -19,6 +19,9 @@ logger = logging.getLogger(__name__) +MAX_BATCH_SIZE = 1024 +"""Where a search stops when its caller sets no bound of its own.""" + NO_GPU_BATCH_SIZE = 8 """Batch size used when there is no GPU to probe.""" diff --git a/learning_loop_node/trainer/cuda.py b/learning_loop_node/trainer/cuda.py index 25039e3b..2d3e2265 100644 --- a/learning_loop_node/trainer/cuda.py +++ b/learning_loop_node/trainer/cuda.py @@ -1,4 +1,4 @@ -"""GPU memory budgeting for trainers that train in-process. +"""GPU memory budgeting and batch-size probing for trainers that train in-process. A trainer that probes for its batch size needs two things a deployment can set: the budget the probe measures against, and a cap that holds the process to it. Both come from one number -- @@ -6,6 +6,14 @@ ``--vram-limit-gb``. It is how one GPU gets shared between processes, and how a training keeps headroom against allocator fragmentation. +The search itself is arithmetic and lives in :mod:`~learning_loop_node.trainer.batch_size`, which +imports no framework. What needs torch is the measuring around it, and that is what this module +adds: hold back a safety margin so a trial competes against a correspondingly smaller card, reset +the peak counters, run the step the node supplies, and tell an out-of-memory failure from a bug. +:func:`probe_batch_size` is the whole probe for a node whose step is one call; a node that has to +build a model first composes :func:`reserve_margin`, :func:`measured_fits` and ``find_batch_size`` +itself. + Unlike the rest of the library this module imports torch, and the package deliberately does not declare it. Capping an allocator is a torch operation with no NVML equivalent, and declaring the dependency would put an ML runtime into a library that is also installed on machines that train @@ -18,11 +26,17 @@ import gc import logging +from collections.abc import Callable import torch +from .batch_size import MAX_BATCH_SIZE, find_batch_size, is_out_of_memory, no_gpu_batch_size + logger = logging.getLogger(__name__) +SAFETY_MARGIN = 0.05 +"""Share of the budget held back while probing, against allocator fragmentation later on.""" + def usable_memory_bytes(vram_limit_gb: float, device: int = 0) -> int: """How much GPU memory this process may allocate, honouring :func:`limit_cuda_memory`. @@ -63,3 +77,83 @@ def free_cuda_memory() -> None: """Release cached CUDA memory, so a following allocation sees the real free space.""" gc.collect() torch.cuda.empty_cache() + + +def probe_batch_size(run_batch: Callable[[int], str | None], *, probe: str = 'batch-size probe', + limit: int = 0, vram_limit_gb: float = 0) -> int: + """Find the largest power-of-two batch size ``run_batch`` fits into. + + For a probe whose measurement is one call -- a forward pass, or a step on a model the caller + already holds. A probe that has to build a throwaway model reserves the margin before + building it, and so composes the pieces this function composes. + + :param run_batch: Runs the batch; may return a detail to append to the log line. + :param probe: Names this probe in the log, so a node running several stays readable. + :param limit: Caps the search; 0 means + :data:`~learning_loop_node.trainer.batch_size.MAX_BATCH_SIZE`. Rounded down to a + power of two. + :param vram_limit_gb: The budget the safety margin is a share of; 0 means the whole card. + :raises InsufficientMemoryError: If not even a batch size of 1 fits. + """ + limit = limit or MAX_BATCH_SIZE + + if not torch.cuda.is_available(): + return no_gpu_batch_size(limit, probe) + + margin = reserve_margin(vram_limit_gb, probe=probe) + chosen = find_batch_size(measured_fits(run_batch, probe=probe), limit=limit) + del margin + free_cuda_memory() + logger.info('%s: selected batch size %d (upper bound %d)', probe, chosen, limit) + return chosen + + +def measured_fits(run_batch: Callable[[int], str | None], *, probe: str, + on_out_of_memory: Callable[[], None] | None = None) -> Callable[[int], bool]: + """Wrap ``run_batch`` into the ``fits`` predicate ``find_batch_size`` searches with. + + The returned predicate frees cached memory before each trial, so a trial does not fail on + what the previous one left behind, and logs the peak the trial reached. + + An out-of-memory failure is the answer "does not fit"; anything else is a bug and is + re-raised. Both arrive as the same exception types, which is why they are told apart by + :func:`~learning_loop_node.trainer.batch_size.is_out_of_memory` rather than by ``except``. + + :param run_batch: Runs the batch; may return a detail to append to the log line. + :param on_out_of_memory: Runs after a trial ran out of memory, to drop what it left behind + (an optimizer's gradients, say). The batch itself is still referenced by the traceback, + so freeing that is left to the next trial. + """ + def fits(batch_size: int) -> bool: + free_cuda_memory() + try: + torch.cuda.reset_peak_memory_stats() + detail = run_batch(batch_size) + torch.cuda.synchronize() + logger.info('%s: %d fits (peak %.2f GB, margin included)%s', probe, batch_size, + torch.cuda.max_memory_allocated() / 1024**3, f'; {detail}' if detail else '') + return True + except (torch.cuda.OutOfMemoryError, RuntimeError, MemoryError) as exc: + if not isinstance(exc, torch.cuda.OutOfMemoryError) and not is_out_of_memory(exc): + raise + logger.info('%s: %d does not fit (%s)', probe, batch_size, type(exc).__name__) + if on_out_of_memory is not None: + on_out_of_memory() + return False + + return fits + + +def reserve_margin(vram_limit_gb: float, *, probe: str) -> torch.Tensor: + """Claim :data:`SAFETY_MARGIN` of the budget, so a trial competes against a smaller card. + + Keep the returned tensor alive for as long as the probe runs: releasing it hands the margin + back, and the size that was chosen against it is no longer the size that was measured. A few + trials cannot show what hours of training do to allocator fragmentation, nor the workspaces + that shapes outside the probe ask cuDNN and cuBLAS for. + + :param vram_limit_gb: The budget the margin is a share of; 0 means the whole card. + """ + margin_bytes = int(usable_memory_bytes(vram_limit_gb) * SAFETY_MARGIN) + logger.info('%s: keeping %.0f MB free as a safety margin', probe, margin_bytes / 1024**2) + return torch.empty(margin_bytes, dtype=torch.uint8, device='cuda') From d9068a973d5a497ec8834a30554e1b74f33f0aca Mon Sep 17 00:00:00 2001 From: Kevin Heye Date: Mon, 7 Sep 2026 14:32:54 +0200 Subject: [PATCH 26/29] Put the callers first, and cut the prose to what the code cannot say A pass over this branch's own code against CONTRIBUTING.md, after reading it back rule by rule. Ordering: cuda.py read bottom-up. probe_batch_size -- which its own docstring calls the whole probe -- sat fourth, below three functions it calls, and usable_memory_bytes sat above its only caller. The order is now the call graph, so a reader meets probe_batch_size first and every callee below all of its callers. Same rule in the unit tests, where _FakeTorch, the load fixture and the small predicate helpers sat above the tests they serve. Defaults: usable_memory_bytes and limit_cuda_memory took device=0, which no caller ever varied. That is a constant hiding in a signature, and the constant was wrong for a trainer that pins itself with torch.cuda.set_device(1) -- it would budget and cap card 0 while training on card 1. Both now let torch resolve the current device. dfine_node passes neither argument, so nothing there needs a companion change. Prose: rules 4 and 7 ask for as few and as short as possible, and no motivation. cuda.py's module docstring goes from 26 lines to 8, keeping the two facts the code cannot state and AGENTS.md does not carry: the cap is a share of the card's total memory rather than of what is free, and torch is imported but undeclared, so only a trainer may import the module. The reasoning about fragmentation, the NVML argument and how the pieces compose already live in AGENTS.md. free_cuda_memory loses its docstring to its own name; is_out_of_memory keeps the one sentence that justifies matching on message text. Eight test docstrings and six comments that restated the test name in other words are gone; the four that state a fact stay -- the bare RuntimeErrors cuDNN and cuBLAS raise, the test whose passing is the absence of a hang, where the confusion-matrix numbers come from, and why HOST looks like the wrong variable name. The Lightning attribution in batch_size.py is untouched: that one is a license, not a style choice. Along the way: test_cuda.py imports relatively like every other test here, its fake torch records capped fractions instead of (-1.0, -1) sentinels, and the new pytest.ini loses the trailing whitespace, the "debbuging" typo and the missing final newline it had inherited from its siblings. 229 unit tests pass. ruff is unchanged at 61 findings across the touched files, all of them in detector_node.py and trainer_logic.py and all predating this branch. Co-Authored-By: Claude Opus 5 --- learning_loop_node/detector/categories.py | 6 +- learning_loop_node/detector/detector_node.py | 3 +- learning_loop_node/detector/geometry.py | 4 +- learning_loop_node/detector/postprocess.py | 10 +- learning_loop_node/helpers/entrypoint.py | 14 +- learning_loop_node/tests/unit/pytest.ini | 8 +- .../tests/unit/test_batch_size.py | 35 ++-- .../tests/unit/test_categories.py | 1 - learning_loop_node/tests/unit/test_cuda.py | 151 ++++++++++-------- .../tests/unit/test_entrypoint.py | 32 ++-- .../tests/unit/test_environment_reader.py | 2 +- .../tests/unit/test_geometry.py | 4 +- learning_loop_node/tests/unit/test_metrics.py | 1 - .../tests/unit/test_postprocess.py | 19 ++- .../tests/unit/test_subprocess.py | 3 +- learning_loop_node/trainer/batch_size.py | 15 +- learning_loop_node/trainer/cuda.py | 134 +++++++--------- learning_loop_node/trainer/metrics.py | 12 +- learning_loop_node/trainer/subprocess.py | 11 +- 19 files changed, 208 insertions(+), 257 deletions(-) diff --git a/learning_loop_node/detector/categories.py b/learning_loop_node/detector/categories.py index a697ce83..67e7d27c 100644 --- a/learning_loop_node/detector/categories.py +++ b/learning_loop_node/detector/categories.py @@ -1,8 +1,4 @@ -"""Resolving a model's categories from what it emits. - -A model reports class indices, or names, and only ``ModelInformation.categories`` gives those -meaning. -""" +"""Resolving the class indices or names a model emits against ``ModelInformation.categories``.""" from ..data_classes import Category, ModelInformation diff --git a/learning_loop_node/detector/detector_node.py b/learning_loop_node/detector/detector_node.py index 3417069e..1b8bf0b7 100644 --- a/learning_loop_node/detector/detector_node.py +++ b/learning_loop_node/detector/detector_node.py @@ -709,8 +709,7 @@ async def upload_images( def add_category_id_to_detections(self, model_info: ModelInformation, image_metadata: ImageMetadata): """Resolve each detection's category id from its name, in metadata a client uploaded. - A name the model does not know is the client's error, not a broken model, so it costs - that one detection its id rather than the whole upload. + A name the model does not know costs that one detection its id, not the whole upload. """ unknown_names: set[str] = set() diff --git a/learning_loop_node/detector/geometry.py b/learning_loop_node/detector/geometry.py index 479362ff..22f9c554 100644 --- a/learning_loop_node/detector/geometry.py +++ b/learning_loop_node/detector/geometry.py @@ -1,7 +1,7 @@ """Box and point clipping shared by every detector node. The loop stores a box as its top-left corner plus a size, which is the form :func:`clip_box` -produces. Centre-anchored model output converts before clipping. +takes and produces. """ @@ -16,8 +16,6 @@ def clip_box( ) -> tuple[int, int, int, int]: """Clip a top-left-anchored box to the image bounds. - :param x1: Left edge of the box. - :param y1: Top edge of the box. :return: The clipped ``(x1, y1, width, height)``; the size is never negative. """ x2 = x1 + width diff --git a/learning_loop_node/detector/postprocess.py b/learning_loop_node/detector/postprocess.py index a6812ddc..7157dc77 100644 --- a/learning_loop_node/detector/postprocess.py +++ b/learning_loop_node/detector/postprocess.py @@ -24,8 +24,8 @@ @dataclass(kw_only=True, slots=True, frozen=True) class Prediction: - """One surviving prediction in the model's coordinates: top-left corner and size in pixels, - unrounded, and the category as an index into :attr:`ModelInformation.categories`.""" + """One surviving prediction in the model's coordinates: unrounded top-left corner and size + in pixels, and the category as an index into :attr:`ModelInformation.categories`.""" x: float y: float @@ -130,11 +130,7 @@ def predictions_from_xyxy( boxes: Sequence[Sequence[float]], scores: Sequence[float], ) -> list[Prediction]: - """Convert xyxy model output into predictions, for models that suppress their own overlaps. - - Corners stay unrounded: :func:`clip_box` rounds once, when the box becomes a - :class:`BoxDetection`. - """ + """Convert xyxy model output into predictions, for models that suppress their own overlaps.""" return [Prediction(x=float(x1), y=float(y1), width=float(x2 - x1), height=float(y2 - y1), category_index=int(label), confidence=float(score)) for label, (x1, y1, x2, y2), score in zip(labels, boxes, scores, strict=True)] diff --git a/learning_loop_node/helpers/entrypoint.py b/learning_loop_node/helpers/entrypoint.py index 5f0d6898..4c42602e 100644 --- a/learning_loop_node/helpers/entrypoint.py +++ b/learning_loop_node/helpers/entrypoint.py @@ -1,8 +1,4 @@ -"""The boilerplate every node's ``main.py`` repeats. - -A node entry point always does the same four things: read a handful of settings, build the -logic object, construct the node, and hand it to uvicorn. Only the middle two are the node's -own, so :func:`node_parser` and :func:`run_node` cover the rest. +"""The boilerplate every node's ``main.py`` repeats: reading the settings, serving the node. Every setting is a flag *and* an environment variable named after it: ``--conf-threshold`` reads ``CONF_THRESHOLD``. ``--host`` and ``--port`` are the exception, reading ``NODE_HOST`` @@ -23,11 +19,9 @@ def node_parser(*, description: str, legacy_env_prefix: str = '') -> configargpa """Build the parser for a node, pre-loaded with the settings every node has. :param legacy_env_prefix: A prefix an earlier version of this node required, e.g. - ``'MY_DETECTOR_'``. Prefixed names are still honoured, with a warning, so a - deployment keeps working until it is updated. Both spellings are accepted: the prefix - on the current environment variable (``MY_DETECTOR_NODE_HOST``) and the prefix on the - flag it was originally applied to (``MY_DETECTOR_HOST``). Leave empty for a node that - has always read unprefixed names. + ``'MY_DETECTOR_'``. Both spellings keep working, with a warning: the prefix on the + current name (``MY_DETECTOR_NODE_HOST``) and the prefix on the flag it was originally + applied to (``MY_DETECTOR_HOST``). Leave empty for a node that never used one. """ parser = _NodeArgumentParser(description=description, legacy_env_prefix=legacy_env_prefix) parser.add_argument('--host', default='0.0.0.0', env_var='NODE_HOST', diff --git a/learning_loop_node/tests/unit/pytest.ini b/learning_loop_node/tests/unit/pytest.ini index 2207d1ac..e3ce9d31 100644 --- a/learning_loop_node/tests/unit/pytest.ini +++ b/learning_loop_node/tests/unit/pytest.ini @@ -2,9 +2,9 @@ python_files = test_*.py asyncio_mode = auto -cache_dir = /tmp/pytest_cache - -# for debbuging tests: +cache_dir = /tmp/pytest_cache + +# for debugging tests: ; log_cli_level = INFO ; log_cli_format = %(asctime)s [%(levelname)8s] %(message)s (%(filename)s:%(lineno)s) -; log_cli_date_format=%Y-%m-%d %H:%M:%S \ No newline at end of file +; log_cli_date_format=%Y-%m-%d %H:%M:%S diff --git a/learning_loop_node/tests/unit/test_batch_size.py b/learning_loop_node/tests/unit/test_batch_size.py index 6b499a6a..9ad3fd3e 100644 --- a/learning_loop_node/tests/unit/test_batch_size.py +++ b/learning_loop_node/tests/unit/test_batch_size.py @@ -13,22 +13,6 @@ ) -def _recording(capacity: int) -> tuple[Callable[[int], bool], list[int]]: - """A `fits` predicate for a machine of `capacity`, plus the sizes it gets asked about.""" - calls: list[int] = [] - - def fits(batch_size: int) -> bool: - calls.append(batch_size) - return batch_size <= capacity - - return fits, calls - - -def _fits_up_to(capacity: int) -> Callable[[int], bool]: - fits, _ = _recording(capacity) - return fits - - def test_the_search_doubles_up_to_the_limit(): fits, calls = _recording(1024) assert find_batch_size(fits, limit=16) == 16 @@ -59,7 +43,6 @@ def test_the_result_is_always_the_largest_power_of_two_that_fits(capacity: int): def test_equal_hardware_yields_an_equal_recipe(): - """Only powers of two, so two machines of similar size train identically.""" assert find_batch_size(_fits_up_to(37), limit=512) == find_batch_size(_fits_up_to(39), limit=512) @@ -90,7 +73,6 @@ def test_the_dataset_limit_stays_usable_for_a_tiny_set(): def test_memory_still_decides_below_the_dataset_limit(): - """The bound is a ceiling only: a card that fits just 2 keeps training at 2.""" assert find_batch_size(_fits_up_to(2), limit=dataset_limit(20)) == 2 @@ -109,10 +91,25 @@ def test_allocation_failures_are_recognised_however_they_surface(message: str): def test_a_real_bug_is_not_mistaken_for_a_full_card(): - """A trainer catching bare RuntimeError treats every crash as 'too big'; this does not.""" assert not is_out_of_memory(RuntimeError('shape mismatch in forward pass')) assert not is_out_of_memory(ValueError('bad config')) def test_the_host_running_out_of_memory_counts_too(): assert is_out_of_memory(MemoryError()) + + +def _recording(capacity: int) -> tuple[Callable[[int], bool], list[int]]: + """A `fits` predicate for a machine of `capacity`, plus the sizes it gets asked about.""" + calls: list[int] = [] + + def fits(batch_size: int) -> bool: + calls.append(batch_size) + return batch_size <= capacity + + return fits, calls + + +def _fits_up_to(capacity: int) -> Callable[[int], bool]: + fits, _ = _recording(capacity) + return fits diff --git a/learning_loop_node/tests/unit/test_categories.py b/learning_loop_node/tests/unit/test_categories.py index 4eacfebb..65492898 100644 --- a/learning_loop_node/tests/unit/test_categories.py +++ b/learning_loop_node/tests/unit/test_categories.py @@ -19,7 +19,6 @@ def test_category_is_resolved_by_index(): @pytest.mark.parametrize('index', [-1, 2, 99]) def test_an_index_outside_the_model_categories_is_an_error(index: int): - # a mismatch between model and metadata must not be silently skipped with pytest.raises(ValueError, match='out of range'): category_by_index(model_information(), index) diff --git a/learning_loop_node/tests/unit/test_cuda.py b/learning_loop_node/tests/unit/test_cuda.py index d7daad0a..c700f3f2 100644 --- a/learning_loop_node/tests/unit/test_cuda.py +++ b/learning_loop_node/tests/unit/test_cuda.py @@ -1,10 +1,10 @@ """Tests for the one library module that imports torch. -There is no torch in the dev extra, and no GPU in CI, so a stand-in is installed under the name -``torch`` before the module is imported. That covers the budgeting arithmetic, the guards and the -search around a probe -- everything this module decides. Whether the cap actually holds, and what -a real step costs, is torch's own business and can only be observed on a card. +There is no torch in the dev extra and no GPU in CI, so a stand-in is installed under the name +``torch`` before the module is imported. Whether the cap actually holds, and what a real step +costs, can only be observed on a card. """ +from __future__ import annotations import importlib import logging @@ -15,66 +15,13 @@ import pytest -from learning_loop_node.trainer.batch_size import MAX_BATCH_SIZE, NO_GPU_BATCH_SIZE -from learning_loop_node.trainer.exceptions import InsufficientMemoryError +from ...trainer.batch_size import MAX_BATCH_SIZE, NO_GPU_BATCH_SIZE +from ...trainer.exceptions import InsufficientMemoryError MODULE = 'learning_loop_node.trainer.cuda' GIB = 1024**3 -class _FakeTorch: - """What the module uses of torch, plus a record of what it asked for.""" - - def __init__(self, *, cuda_available: bool, total_gb: float, peak_gb: float) -> None: - self.capped: list[tuple[float, int]] = [] - self.allocated: list[int] = [] - - class OutOfMemoryError(RuntimeError): - """Torch's own; note the module may not rely on its message.""" - - self.OutOfMemoryError = OutOfMemoryError # named as torch spells it - self.module = types.ModuleType('torch') - self.module.uint8 = 'uint8' # type: ignore[attr-defined] - self.module.empty = self._empty # type: ignore[attr-defined] - self.module.cuda = types.SimpleNamespace( # type: ignore[attr-defined] - OutOfMemoryError=OutOfMemoryError, - is_available=lambda: cuda_available, - empty_cache=lambda: self.capped.append((-1.0, -1)), - get_device_properties=lambda device: types.SimpleNamespace(total_memory=int(total_gb * GIB)), - set_per_process_memory_fraction=lambda fraction, device: self.capped.append((fraction, device)), - reset_peak_memory_stats=lambda: None, - max_memory_allocated=lambda: int(peak_gb * GIB), - synchronize=lambda: None, - ) - - def _empty(self, count: int, dtype: str, device: str) -> object: - assert (dtype, device) == ('uint8', 'cuda') - self.allocated.append(count) - return object() - - -@pytest.fixture(name='load') -def load_fixture(monkeypatch: pytest.MonkeyPatch) -> Callable[..., tuple[Any, _FakeTorch]]: - """Import the module against a fake torch; returns it and the stand-in it ran against.""" - def load(*, cuda_available: bool = True, total_gb: float = 8.0, peak_gb: float = 2.0): - fake = _FakeTorch(cuda_available=cuda_available, total_gb=total_gb, peak_gb=peak_gb) - monkeypatch.setitem(sys.modules, 'torch', fake.module) - monkeypatch.delitem(sys.modules, MODULE, raising=False) - return importlib.import_module(MODULE), fake - - yield load - sys.modules.pop(MODULE, None) - - -def _fits_up_to(largest: int, fake: _FakeTorch, ran: list[int]) -> Callable[[int], None]: - """A step that runs out of memory above ``largest``, recording every size it was asked for.""" - def run_batch(batch_size: int) -> None: - ran.append(batch_size) - if batch_size > largest: - raise fake.OutOfMemoryError('tried to allocate 20.00 GiB') - return run_batch - - # --- the budget and the cap --- def test_no_limit_means_the_whole_card(load): @@ -89,7 +36,6 @@ def test_a_limit_below_the_card_is_the_budget(load): def test_a_limit_above_the_card_is_clamped_to_it(load): - # otherwise the safety margin would be a share of memory that does not exist cuda, _ = load(total_gb=8.0) assert cuda.usable_memory_bytes(16) == 8 * GIB @@ -97,7 +43,7 @@ def test_a_limit_above_the_card_is_clamped_to_it(load): def test_the_cap_is_the_limits_share_of_the_card(load): cuda, fake = load(total_gb=8.0) cuda.limit_cuda_memory(2) - assert fake.capped == [(0.25, 0)] + assert fake.capped == [0.25] def test_no_limit_caps_nothing(load): @@ -114,7 +60,6 @@ def test_nothing_is_capped_without_cuda(load): def test_a_limit_the_card_cannot_reach_warns_instead_of_capping(load, caplog): - # capping at a fraction >= 1 would be a no-op that reads as a limit having been applied cuda, fake = load(total_gb=8.0) with caplog.at_level(logging.WARNING): cuda.limit_cuda_memory(8) @@ -122,10 +67,17 @@ def test_a_limit_the_card_cannot_reach_warns_instead_of_capping(load, caplog): assert 'exceeds the card capacity' in caplog.text +def test_the_budget_and_the_cap_follow_the_current_device(load): + cuda, fake = load(total_gb=8.0) + cuda.usable_memory_bytes(2) + cuda.limit_cuda_memory(2) + assert fake.asked_devices and all(device is None for device in fake.asked_devices) + + def test_freeing_empties_the_cache(load): cuda, fake = load() cuda.free_cuda_memory() - assert fake.capped == [(-1.0, -1)] + assert fake.cache_clears == 1 # --- the safety margin --- @@ -152,7 +104,6 @@ def test_the_probe_keeps_the_largest_size_that_fits(load): def test_the_probe_reserves_the_margin_before_it_measures(load): - # a trial that ran against the full card would choose a size the training cannot keep cuda, fake = load(total_gb=8.0) def run_batch(_: int) -> None: @@ -174,7 +125,6 @@ def test_an_unset_limit_stops_at_the_maximum(load): def test_a_probe_without_a_gpu_does_not_run_the_step(load): - """With no card to measure, the fallback caps the batch and never runs the step.""" cuda, fake = load(cuda_available=False) ran: list[int] = [] assert cuda.probe_batch_size(_fits_up_to(1024, fake, ran), limit=32) == NO_GPU_BATCH_SIZE @@ -222,7 +172,6 @@ def run_batch(batch_size: int) -> None: def test_a_failure_that_is_not_about_memory_is_a_bug_and_propagates(load): - # a probe that swallows this would report the smallest batch size as the card's fault cuda, _ = load() def run_batch(batch_size: int) -> None: @@ -249,3 +198,73 @@ def test_only_a_trial_that_ran_out_of_memory_gets_cleaned_up_after(load): assert [fits(size) for size in (1, 2, 4)] == [True, True, False] assert cleaned == [3], 'once, after the third trial' + + +# --- the card that is not there --- + +@pytest.fixture(name='load') +def load_fixture(monkeypatch: pytest.MonkeyPatch) -> Callable[..., tuple[Any, _FakeTorch]]: + """Import the module against a fake torch; returns it and the stand-in it ran against.""" + def load(*, cuda_available: bool = True, total_gb: float = 8.0, peak_gb: float = 2.0): + fake = _FakeTorch(cuda_available=cuda_available, total_gb=total_gb, peak_gb=peak_gb) + monkeypatch.setitem(sys.modules, 'torch', fake.module) + monkeypatch.delitem(sys.modules, MODULE, raising=False) + return importlib.import_module(MODULE), fake + + yield load + sys.modules.pop(MODULE, None) + + +def _fits_up_to(largest: int, fake: _FakeTorch, ran: list[int]) -> Callable[[int], None]: + """A step that runs out of memory above ``largest``, recording every size it was asked for.""" + def run_batch(batch_size: int) -> None: + ran.append(batch_size) + if batch_size > largest: + raise fake.OutOfMemoryError('tried to allocate 20.00 GiB') + return run_batch + + +class _FakeTorch: + """What the module uses of torch, plus a record of what it asked for.""" + + def __init__(self, *, cuda_available: bool, total_gb: float, peak_gb: float) -> None: + self.capped: list[float] = [] + self.allocated: list[int] = [] + self.asked_devices: list[int | None] = [] + self.cache_clears = 0 + + class OutOfMemoryError(RuntimeError): + """Torch's own; note the module may not rely on its message.""" + + self.OutOfMemoryError = OutOfMemoryError # named as torch spells it + self.module = types.ModuleType('torch') + self.module.uint8 = 'uint8' # type: ignore[attr-defined] + self.module.empty = self._empty # type: ignore[attr-defined] + self.module.cuda = types.SimpleNamespace( # type: ignore[attr-defined] + OutOfMemoryError=OutOfMemoryError, + is_available=lambda: cuda_available, + empty_cache=self._empty_cache, + get_device_properties=self._device_properties(total_gb), + set_per_process_memory_fraction=self._cap, + reset_peak_memory_stats=lambda: None, + max_memory_allocated=lambda: int(peak_gb * GIB), + synchronize=lambda: None, + ) + + def _empty(self, count: int, dtype: str, device: str) -> object: + assert (dtype, device) == ('uint8', 'cuda') + self.allocated.append(count) + return object() + + def _empty_cache(self) -> None: + self.cache_clears += 1 + + def _device_properties(self, total_gb: float) -> Callable[[int | None], Any]: + def get_device_properties(device: int | None = None): + self.asked_devices.append(device) + return types.SimpleNamespace(total_memory=int(total_gb * GIB)) + return get_device_properties + + def _cap(self, fraction: float, device: int | None = None) -> None: + self.asked_devices.append(device) + self.capped.append(fraction) diff --git a/learning_loop_node/tests/unit/test_entrypoint.py b/learning_loop_node/tests/unit/test_entrypoint.py index 03a29f6b..9eed399a 100644 --- a/learning_loop_node/tests/unit/test_entrypoint.py +++ b/learning_loop_node/tests/unit/test_entrypoint.py @@ -6,19 +6,6 @@ 'MY_DETECTOR_HOST', 'MY_DETECTOR_PORT', 'MY_DETECTOR_NODE_HOST') -@pytest.fixture(autouse=True) -def clean_env(monkeypatch: pytest.MonkeyPatch): - """Every test starts without the variables it is about to set.""" - for name in MANAGED: - monkeypatch.delenv(name, raising=False) - - -def _parser(**kwargs): - parser = node_parser(description='a node', **kwargs) - parser.add_argument('--weight-type', default='FP16') - return parser - - def test_every_node_gets_a_host_and_a_port(): args = _parser().parse_args([]) assert (args.host, args.port) == ('0.0.0.0', 80) @@ -35,7 +22,6 @@ def test_a_setting_is_read_from_the_variable_named_after_its_flag(monkeypatch: p def test_the_loop_own_host_is_never_mistaken_for_the_bind_address(monkeypatch: pytest.MonkeyPatch): - """HOST is the loop's address. Binding uvicorn to it would leave the node unreachable.""" monkeypatch.setenv('HOST', 'preview.learning-loop.ai') assert _parser().parse_args([]).host == '0.0.0.0' @@ -74,8 +60,7 @@ def test_a_node_without_a_legacy_prefix_ignores_prefixed_names(monkeypatch: pyte def test_the_bind_address_is_still_read_under_the_name_the_prefix_gave_the_flag( monkeypatch: pytest.MonkeyPatch): - """A prefix used to be applied to the flag, so `--host` was `HOST`, not - `NODE_HOST`. That is the only spelling a deployment can have set.""" + """A prefix used to be applied to the flag, so `--host` was `HOST`.""" monkeypatch.setenv('MY_DETECTOR_HOST', '127.0.0.1') monkeypatch.setenv('MY_DETECTOR_PORT', '8099') args = _parser(legacy_env_prefix='MY_DETECTOR_').parse_args([]) @@ -91,7 +76,6 @@ def test_the_renamed_setting_warns_which_name_to_use_instead(monkeypatch: pytest def test_the_prefixed_current_name_is_honoured_too(monkeypatch: pytest.MonkeyPatch): - """A node whose prefixed spelling was already the current one keeps working.""" monkeypatch.setenv('MY_DETECTOR_NODE_HOST', '10.0.0.1') assert _parser(legacy_env_prefix='MY_DETECTOR_').parse_args([]).host == '10.0.0.1' @@ -104,6 +88,18 @@ def test_the_current_name_wins_over_both_prefixed_spellings(monkeypatch: pytest. def test_the_loop_own_host_is_not_adopted_by_a_prefixed_node(monkeypatch: pytest.MonkeyPatch): - """`HOST` is the loop's address even here: only the *prefixed* spelling is a bind address.""" monkeypatch.setenv('HOST', 'preview.learning-loop.ai') assert _parser(legacy_env_prefix='MY_DETECTOR_').parse_args([]).host == '0.0.0.0' + + +@pytest.fixture(autouse=True) +def clean_env(monkeypatch: pytest.MonkeyPatch): + """Every test starts without the variables it is about to set.""" + for name in MANAGED: + monkeypatch.delenv(name, raising=False) + + +def _parser(**kwargs): + parser = node_parser(description='a node', **kwargs) + parser.add_argument('--weight-type', default='FP16') + return parser diff --git a/learning_loop_node/tests/unit/test_environment_reader.py b/learning_loop_node/tests/unit/test_environment_reader.py index a66a844e..df957bc0 100644 --- a/learning_loop_node/tests/unit/test_environment_reader.py +++ b/learning_loop_node/tests/unit/test_environment_reader.py @@ -28,7 +28,7 @@ def test_either_name_alone_is_read(monkeypatch: pytest.MonkeyPatch): def test_a_disagreement_resolves_to_the_preferred_name(monkeypatch: pytest.MonkeyPatch): - """Returning nothing here would let host() fall back to its default, which is production.""" + """Returning nothing here would let host() fall back to its default.""" monkeypatch.setenv('LOOP_HOST', 'preview.learning-loop.ai') monkeypatch.setenv('HOST', 'learning-loop.ai') assert environment_reader.host(default='learning-loop.ai') == 'preview.learning-loop.ai' diff --git a/learning_loop_node/tests/unit/test_geometry.py b/learning_loop_node/tests/unit/test_geometry.py index a2f3836b..76fa2c0b 100644 --- a/learning_loop_node/tests/unit/test_geometry.py +++ b/learning_loop_node/tests/unit/test_geometry.py @@ -11,9 +11,7 @@ def test_box_is_clipped_to_the_image_bounds(): def test_box_fully_outside_the_image_collapses_to_zero_size(): - # The corner is only clamped at the lower bound, so it stays at 200 — the zero size is - # what marks the box as empty. Detector output cannot reach here, because - # non_max_suppression already clips every box into the image. + # the corner is clamped at the lower bound only; the zero size is what marks the box empty assert clip_box(x1=200, y1=200, width=10, height=10, img_width=100, img_height=100) == (200, 200, 0, 0) diff --git a/learning_loop_node/tests/unit/test_metrics.py b/learning_loop_node/tests/unit/test_metrics.py index 7c8537dc..b9a0fe6b 100644 --- a/learning_loop_node/tests/unit/test_metrics.py +++ b/learning_loop_node/tests/unit/test_metrics.py @@ -11,7 +11,6 @@ def test_the_score_averages_the_categories_instead_of_pooling_them(): def test_a_rare_category_carries_the_same_weight(): - """A category the model never finds halves the score, however few instances it has.""" assert macro_f1({'frequent': {'tp': 1000, 'fp': 0, 'fn': 0}, 'rare': {'tp': 0, 'fp': 0, 'fn': 3}}) == pytest.approx(0.5) diff --git a/learning_loop_node/tests/unit/test_postprocess.py b/learning_loop_node/tests/unit/test_postprocess.py index 6d3ddcf4..93cb6b48 100644 --- a/learning_loop_node/tests/unit/test_postprocess.py +++ b/learning_loop_node/tests/unit/test_postprocess.py @@ -17,12 +17,7 @@ POINT = Category(id='uuid-point', name='weed', type=CategoryType.Point) -def model_information(*categories: Category) -> ModelInformation: - return ModelInformation(id='model-uuid', host='localhost', organization='zauberzeug', - project='pytest', version='1.2', categories=list(categories or (BOX, POINT))) - - -# ---------------------------------------------------------------- iou and suppression +# --- iou and suppression --- def test_identical_boxes_have_an_iou_of_one(): box = np.array([[0, 0, 10, 10]], dtype=np.float32) @@ -60,7 +55,7 @@ def test_suppression_clips_boxes_into_the_image(): assert list(kept_boxes[0]) == [0, 0, 99, 99] -# ---------------------------------------------------------------- post_process +# --- post_process --- def test_post_process_drops_predictions_below_the_confidence_threshold(): boxes = np.array([[10, 10, 60, 60], [100, 100, 150, 150]], dtype=np.float32) @@ -70,7 +65,7 @@ def test_post_process_drops_predictions_below_the_confidence_threshold(): prediction = result[0] assert (prediction.x, prediction.y, prediction.width, prediction.height) == (10, 10, 50, 50) assert prediction.category_index == 0 - # the model's own float32 score, no longer rounded to two decimals on the way out + # the model's own float32 score, not rounded assert prediction.confidence == pytest.approx(0.9) @@ -90,7 +85,7 @@ def test_converting_already_suppressed_output_requires_matching_lengths(): predictions_from_xyxy(labels=[1.0, 2.0], boxes=[[0.0, 0.0, 1.0, 1.0]], scores=[0.5]) -# ---------------------------------------------------------------- building the containers +# --- building the containers --- def test_a_box_category_becomes_a_box_detection(): metadata = to_image_metadata([Prediction(x=10, y=20, width=30, height=40, category_index=0, confidence=0.9)], model_information(), 200, 200) @@ -137,7 +132,6 @@ def test_the_trainer_container_carries_the_image_id(): def test_trainer_and_detector_paths_agree_on_the_same_detections(): - """The whole point of sharing this code: auto-detections and live detections must match.""" predictions = [Prediction(x=-5, y=-5, width=60, height=60, category_index=0, confidence=0.9), Prediction(x=100, y=100, width=40, height=40, category_index=1, confidence=0.7), Prediction(x=5, y=5, width=1, height=1, category_index=0, confidence=0.5)] metadata = to_image_metadata(predictions, model_information(), 200, 200) @@ -147,3 +141,8 @@ def test_trainer_and_detector_paths_agree_on_the_same_detections(): [(d.x, d.y, d.width, d.height, d.category_id) for d in metadata.box_detections] assert [(d.x, d.y, d.category_id) for d in result.point_detections] == \ [(d.x, d.y, d.category_id) for d in metadata.point_detections] + + +def model_information(*categories: Category) -> ModelInformation: + return ModelInformation(id='model-uuid', host='localhost', organization='zauberzeug', + project='pytest', version='1.2', categories=list(categories or (BOX, POINT))) diff --git a/learning_loop_node/tests/unit/test_subprocess.py b/learning_loop_node/tests/unit/test_subprocess.py index f2517ced..b197cc9b 100644 --- a/learning_loop_node/tests/unit/test_subprocess.py +++ b/learning_loop_node/tests/unit/test_subprocess.py @@ -31,7 +31,6 @@ async def test_a_generator_that_yields_nothing_simply_finishes(): async def test_a_failure_in_the_process_is_raised_in_the_caller(): - """Otherwise a crashed training would look like a training that finished.""" with pytest.raises(RuntimeError, match='the training crashed'): async with iterator_cpu_bound(raising_after, 2) as iterator: await _collect(iterator) @@ -51,4 +50,4 @@ async def test_leaving_early_does_not_leave_the_process_running(): async for item in iterator: if item == 2: break - # the context manager killed and joined the process; reaching here without hanging is the test + # reaching here without hanging is the test diff --git a/learning_loop_node/trainer/batch_size.py b/learning_loop_node/trainer/batch_size.py index 78b46f9c..8fb0418c 100644 --- a/learning_loop_node/trainer/batch_size.py +++ b/learning_loop_node/trainer/batch_size.py @@ -1,11 +1,8 @@ """Choosing a batch size by probing, rather than configuring one. -The largest batch that fits depends on the model, the image resolution and the card. Probing -runs a representative step at doubling sizes and keeps the last one that survived; what "a -representative step" means is the trainer's business — it supplies a ``fits`` predicate. -Nothing here imports a deep-learning framework. - -Only powers of two are visited, so equal hardware yields an equal recipe. +The trainer supplies a ``fits`` predicate that runs a representative step at doubling sizes. +Only powers of two are visited, so equal hardware yields an equal recipe. Nothing here imports a +deep-learning framework. Adapted from PyTorch Lightning's ``BatchSizeFinder`` (power-scaling mode). Copyright The Lightning AI team. Licensed under the Apache License, Version 2.0. @@ -75,10 +72,8 @@ def batch_count(sample_count: int, batch_size: int) -> int: def is_out_of_memory(exception: BaseException) -> bool: """Whether the exception signals exhausted memory, on the GPU or the host. - Allocation failures do not all surface as a framework's dedicated error type: cuDNN and - cuBLAS workspaces raise a plain ``RuntimeError``. Catching those by message is what keeps a - probe from mistaking a real bug for a full card — a trainer that catches bare - ``RuntimeError`` around its probe silently treats every crash as "too big". + cuDNN and cuBLAS workspace failures raise a plain ``RuntimeError``, so the message has to be + matched too. """ if isinstance(exception, MemoryError): return True diff --git a/learning_loop_node/trainer/cuda.py b/learning_loop_node/trainer/cuda.py index 2d3e2265..45815d33 100644 --- a/learning_loop_node/trainer/cuda.py +++ b/learning_loop_node/trainer/cuda.py @@ -1,26 +1,11 @@ """GPU memory budgeting and batch-size probing for trainers that train in-process. -A trainer that probes for its batch size needs two things a deployment can set: the budget the -probe measures against, and a cap that holds the process to it. Both come from one number -- -gigabytes of the card this training may use -- which is what a node exposes as -``--vram-limit-gb``. It is how one GPU gets shared between processes, and how a training keeps -headroom against allocator fragmentation. - -The search itself is arithmetic and lives in :mod:`~learning_loop_node.trainer.batch_size`, which -imports no framework. What needs torch is the measuring around it, and that is what this module -adds: hold back a safety margin so a trial competes against a correspondingly smaller card, reset -the peak counters, run the step the node supplies, and tell an out-of-memory failure from a bug. -:func:`probe_batch_size` is the whole probe for a node whose step is one call; a node that has to -build a model first composes :func:`reserve_margin`, :func:`measured_fits` and ``find_batch_size`` -itself. - -Unlike the rest of the library this module imports torch, and the package deliberately does not -declare it. Capping an allocator is a torch operation with no NVML equivalent, and declaring the -dependency would put an ML runtime into a library that is also installed on machines that train -nothing. Only a trainer imports this module, and a trainer brings torch already. - -Note the cap is relative to the card's *total* memory, not to what is free, so it does not -protect against another process claiming memory first. +``--vram-limit-gb`` -- gigabytes of the card this training may use -- becomes both the budget a +probe measures against and the cap that holds the process to it, on whichever GPU the calling +process is using. The cap is a share of the card's *total* memory, not of what is free. + +This is the one module in the library that imports torch, which the package does not declare, so +only a trainer may import it. The search itself is in :mod:`~learning_loop_node.trainer.batch_size`. """ from __future__ import annotations @@ -38,60 +23,17 @@ """Share of the budget held back while probing, against allocator fragmentation later on.""" -def usable_memory_bytes(vram_limit_gb: float, device: int = 0) -> int: - """How much GPU memory this process may allocate, honouring :func:`limit_cuda_memory`. - - :param vram_limit_gb: 0 or less means the whole card. - """ - total_bytes = torch.cuda.get_device_properties(device).total_memory - if vram_limit_gb <= 0: - return total_bytes - return min(total_bytes, int(vram_limit_gb * 1024**3)) - - -def limit_cuda_memory(vram_limit_gb: float, device: int = 0) -> None: - """Cap how much of the GPU this process may allocate, to ``vram_limit_gb`` gigabytes. - - Call this once per process that touches the GPU -- a spawned training process included, - since the cap does not survive the spawn. - - :param vram_limit_gb: 0 or less means no cap. - """ - if vram_limit_gb <= 0 or not torch.cuda.is_available(): - return - - total_bytes = torch.cuda.get_device_properties(device).total_memory - fraction = vram_limit_gb * 1024**3 / total_bytes - total_gb = total_bytes / 1024**3 - - if fraction >= 1.0: - logger.warning('VRAM limit of %.1f GB exceeds the card capacity of %.1f GB; not limiting', - vram_limit_gb, total_gb) - return - - torch.cuda.set_per_process_memory_fraction(fraction, device) - logger.info('Limiting VRAM usage to %.1f GB of %.1f GB (%.0f%%)', vram_limit_gb, total_gb, fraction * 100) - - -def free_cuda_memory() -> None: - """Release cached CUDA memory, so a following allocation sees the real free space.""" - gc.collect() - torch.cuda.empty_cache() - - def probe_batch_size(run_batch: Callable[[int], str | None], *, probe: str = 'batch-size probe', limit: int = 0, vram_limit_gb: float = 0) -> int: """Find the largest power-of-two batch size ``run_batch`` fits into. - For a probe whose measurement is one call -- a forward pass, or a step on a model the caller - already holds. A probe that has to build a throwaway model reserves the margin before - building it, and so composes the pieces this function composes. + For a probe whose measurement is one call. A probe that has to build a throwaway model first + composes :func:`reserve_margin`, :func:`measured_fits` and ``find_batch_size`` itself. :param run_batch: Runs the batch; may return a detail to append to the log line. :param probe: Names this probe in the log, so a node running several stays readable. - :param limit: Caps the search; 0 means - :data:`~learning_loop_node.trainer.batch_size.MAX_BATCH_SIZE`. Rounded down to a - power of two. + :param limit: Caps the search, rounded down to a power of two; 0 means + :data:`~learning_loop_node.trainer.batch_size.MAX_BATCH_SIZE`. :param vram_limit_gb: The budget the safety margin is a share of; 0 means the whole card. :raises InsufficientMemoryError: If not even a batch size of 1 fits. """ @@ -112,17 +54,13 @@ def measured_fits(run_batch: Callable[[int], str | None], *, probe: str, on_out_of_memory: Callable[[], None] | None = None) -> Callable[[int], bool]: """Wrap ``run_batch`` into the ``fits`` predicate ``find_batch_size`` searches with. - The returned predicate frees cached memory before each trial, so a trial does not fail on - what the previous one left behind, and logs the peak the trial reached. - - An out-of-memory failure is the answer "does not fit"; anything else is a bug and is - re-raised. Both arrive as the same exception types, which is why they are told apart by - :func:`~learning_loop_node.trainer.batch_size.is_out_of_memory` rather than by ``except``. + An out-of-memory failure is the answer "does not fit"; anything else is re-raised. Both + arrive as the same exception types, so they are told apart by + :func:`~learning_loop_node.trainer.batch_size.is_out_of_memory`, not by ``except``. :param run_batch: Runs the batch; may return a detail to append to the log line. :param on_out_of_memory: Runs after a trial ran out of memory, to drop what it left behind - (an optimizer's gradients, say). The batch itself is still referenced by the traceback, - so freeing that is left to the next trial. + (an optimizer's gradients, say). """ def fits(batch_size: int) -> bool: free_cuda_memory() @@ -148,12 +86,50 @@ def reserve_margin(vram_limit_gb: float, *, probe: str) -> torch.Tensor: """Claim :data:`SAFETY_MARGIN` of the budget, so a trial competes against a smaller card. Keep the returned tensor alive for as long as the probe runs: releasing it hands the margin - back, and the size that was chosen against it is no longer the size that was measured. A few - trials cannot show what hours of training do to allocator fragmentation, nor the workspaces - that shapes outside the probe ask cuDNN and cuBLAS for. + back, and the chosen size is no longer the size that was measured. :param vram_limit_gb: The budget the margin is a share of; 0 means the whole card. """ margin_bytes = int(usable_memory_bytes(vram_limit_gb) * SAFETY_MARGIN) logger.info('%s: keeping %.0f MB free as a safety margin', probe, margin_bytes / 1024**2) return torch.empty(margin_bytes, dtype=torch.uint8, device='cuda') + + +def usable_memory_bytes(vram_limit_gb: float) -> int: + """How much GPU memory this process may allocate, honouring :func:`limit_cuda_memory`. + + :param vram_limit_gb: 0 or less means the whole card. + """ + total_bytes = torch.cuda.get_device_properties(None).total_memory + if vram_limit_gb <= 0: + return total_bytes + return min(total_bytes, int(vram_limit_gb * 1024**3)) + + +def limit_cuda_memory(vram_limit_gb: float) -> None: + """Cap how much of the GPU this process may allocate, to ``vram_limit_gb`` gigabytes. + + Call this once per process that touches the GPU, a spawned training process included: the + cap does not survive the spawn. + + :param vram_limit_gb: 0 or less means no cap. + """ + if vram_limit_gb <= 0 or not torch.cuda.is_available(): + return + + total_bytes = torch.cuda.get_device_properties(None).total_memory + fraction = vram_limit_gb * 1024**3 / total_bytes + total_gb = total_bytes / 1024**3 + + if fraction >= 1.0: + logger.warning('VRAM limit of %.1f GB exceeds the card capacity of %.1f GB; not limiting', + vram_limit_gb, total_gb) + return + + torch.cuda.set_per_process_memory_fraction(fraction, None) + logger.info('Limiting VRAM usage to %.1f GB of %.1f GB (%.0f%%)', vram_limit_gb, total_gb, fraction * 100) + + +def free_cuda_memory() -> None: + gc.collect() + torch.cuda.empty_cache() diff --git a/learning_loop_node/trainer/metrics.py b/learning_loop_node/trainer/metrics.py index 6fb6baf9..0ce6c43d 100644 --- a/learning_loop_node/trainer/metrics.py +++ b/learning_loop_node/trainer/metrics.py @@ -1,18 +1,12 @@ -"""Scoring a training from the confusion matrix the loop stores. - -The loop keeps one ``{'tp': .., 'fp': .., 'fn': ..}`` per category, which is what -``TrainerLogicGeneric._get_new_best_training_state`` returns. -""" +"""Scoring a training from the confusion matrix the loop stores: one +``{'tp': .., 'fp': .., 'fn': ..}`` per category.""" import statistics from collections.abc import Mapping def macro_f1(confusion_matrix: Mapping[str, Mapping[str, int]]) -> float: - """The unweighted mean of the per-category F1 scores, as the loop's UI shows it by default. - - Unweighted, so a rare category weighs as much as a frequent one. - """ + """The unweighted mean of the per-category F1 scores, as the loop's UI shows it by default.""" scores = [category_f1(counts) for counts in confusion_matrix.values()] return statistics.mean(scores) if scores else 0.0 diff --git a/learning_loop_node/trainer/subprocess.py b/learning_loop_node/trainer/subprocess.py index 05386064..118fbf5d 100644 --- a/learning_loop_node/trainer/subprocess.py +++ b/learning_loop_node/trainer/subprocess.py @@ -1,12 +1,9 @@ """Run a blocking, CPU-bound generator in its own process without blocking the event loop. -:func:`iterator_cpu_bound` runs the generator in a spawned process and yields what it produces -through a ``maxsize=1`` queue, so the producer never runs more than one item ahead of the -consumer. The context is spawn on every platform, not the local default, so a process that has -already initialised CUDA is never forked; ``it`` and its arguments must therefore be picklable. - -Exceptions raised inside the process are re-raised in the caller, and the process is killed if -the caller leaves the context early. +The queue has ``maxsize=1``, so the producer never runs more than one item ahead of the consumer. +The context is spawn on every platform, so ``it`` and its arguments must be picklable and a +process that already initialised CUDA is never forked. An exception raised inside the process is +re-raised in the caller, and the process is killed if the caller leaves the context early. """ from __future__ import annotations From 6ad79cf4d272081a6e2343c45dd167ea835bee9e Mon Sep 17 00:00:00 2001 From: Kevin Heye Date: Tue, 8 Sep 2026 14:08:07 +0200 Subject: [PATCH 27/29] Run the slow suites on the version the nodes actually ship The matrix tested the two ends of requires-python, 3.10 and 3.13, and so tested neither version a node runs: dfine_node's and yolov5_node's trainers are both python:3.12.11-slim-bookworm, and the dfine detector runs on ubuntu24.04, whose system python is 3.12. The suites that talk to a real loop are exactly the ones where that gap matters, so 3.13 gives way to 3.12 there. 3.13 keeps its coverage in the unit job, which runs the whole 3.10 through 3.13 range and gates this one -- and no node runs it in production yet. The floor stays, since requires-python promises it. Co-Authored-By: Claude Opus 5 --- .github/workflows/pytest.yml | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/.github/workflows/pytest.yml b/.github/workflows/pytest.yml index 929eb642..958bd0a0 100644 --- a/.github/workflows/pytest.yml +++ b/.github/workflows/pytest.yml @@ -34,9 +34,9 @@ jobs: # would delete each other's project max-parallel: 1 matrix: - # the ends of what requires-python promises; the unit job covers 3.11 and 3.12 too, - # and no node runs anything outside this range - python-version: ["3.10", "3.13"] + # the floor requires-python promises, and the version every node ships with; the unit + # job covers 3.11 and 3.13 too + python-version: ["3.10", "3.12"] steps: - uses: actions/checkout@v4 - uses: astral-sh/setup-uv@v3 From a01534e40a4f2c2e8f536a9e9235c7d9cfd12250 Mon Sep 17 00:00:00 2001 From: jan Date: Tue, 8 Sep 2026 15:11:08 +0200 Subject: [PATCH 28/29] Release the probe's margin on every exit, and log the bound the search used When find_batch_size raised, the margin tensor stayed referenced by the traceback until the caller's except block ended, so the reserved share was still held while a fallback ran. The release now sits in a finally. The limit is rounded down to a power of two once, before the search; the log line printed the raw value, so a limit of 12 read as if 12 had been tried. --- learning_loop_node/trainer/cuda.py | 12 +++++++----- 1 file changed, 7 insertions(+), 5 deletions(-) diff --git a/learning_loop_node/trainer/cuda.py b/learning_loop_node/trainer/cuda.py index 45815d33..1b126bb0 100644 --- a/learning_loop_node/trainer/cuda.py +++ b/learning_loop_node/trainer/cuda.py @@ -15,7 +15,7 @@ import torch -from .batch_size import MAX_BATCH_SIZE, find_batch_size, is_out_of_memory, no_gpu_batch_size +from .batch_size import MAX_BATCH_SIZE, find_batch_size, is_out_of_memory, no_gpu_batch_size, smaller_pot logger = logging.getLogger(__name__) @@ -37,15 +37,17 @@ def probe_batch_size(run_batch: Callable[[int], str | None], *, probe: str = 'ba :param vram_limit_gb: The budget the safety margin is a share of; 0 means the whole card. :raises InsufficientMemoryError: If not even a batch size of 1 fits. """ - limit = limit or MAX_BATCH_SIZE + limit = smaller_pot(limit or MAX_BATCH_SIZE) if not torch.cuda.is_available(): return no_gpu_batch_size(limit, probe) margin = reserve_margin(vram_limit_gb, probe=probe) - chosen = find_batch_size(measured_fits(run_batch, probe=probe), limit=limit) - del margin - free_cuda_memory() + try: + chosen = find_batch_size(measured_fits(run_batch, probe=probe), limit=limit) + finally: + del margin + free_cuda_memory() logger.info('%s: selected batch size %d (upper bound %d)', probe, chosen, limit) return chosen From ada6161ac822f25f4a0aee62c441ef34a79e2813 Mon Sep 17 00:00:00 2001 From: Kevin Heye Date: Wed, 9 Sep 2026 11:08:08 +0200 Subject: [PATCH 29/29] Gate the branch on one check name the matrix cannot change The ruleset requires pytest_3_10 and pytest_3_13, which were job ids until the two jobs became one matrix job. A matrix job's check run is named " ()", so both contexts stopped existing -- and an unreported required check is not a failure but "expected", which waits forever. Every PR carrying the refactor is stuck on a status that nothing can report, this one included. all-green needs both jobs and passes only when both succeeded, so the ruleset can name it instead and stop tracking the matrix. always() plus an explicit test of each result is what makes it a gate: a job that is skipped because its dependency failed reports no status at all, and a required check that never reports blocks rather than fails. Swapping the two contexts for all-green in the ruleset needs a repo admin; until then this changes nothing about what CI runs. Co-Authored-By: Claude Opus 5 --- .github/workflows/pytest.yml | 17 +++++++++++++++++ 1 file changed, 17 insertions(+) diff --git a/.github/workflows/pytest.yml b/.github/workflows/pytest.yml index 958bd0a0..041c1958 100644 --- a/.github/workflows/pytest.yml +++ b/.github/workflows/pytest.yml @@ -83,6 +83,23 @@ jobs: LOOP_PASSWORD: ${{ secrets.LEARNING_LOOP_ADMIN_PASSWORD }} run: uv run --no-sync python -m pytest mock_trainer -v + all-green: + # the one status check the branch ruleset requires: a matrix job's check run is named + # " ()", so requiring the entries themselves ties the ruleset to the matrix + # and every change to it leaves the ruleset waiting for a name nothing reports any more + needs: + - unit + - pytest + if: always() # a required check that is skipped reports nothing and waits forever + runs-on: ubuntu-latest + steps: + - name: check the jobs it gates + # skipped and cancelled have to fail too -- only success is green + run: | + echo "unit: ${{ needs.unit.result }}, pytest: ${{ needs.pytest.result }}" + [ "${{ needs.unit.result }}" = success ] || exit 1 + [ "${{ needs.pytest.result }}" = success ] || exit 1 + slack: needs: - unit