learning-loop-node 0.22.0__py3-none-any.whl → 0.23.0__py3-none-any.whl

learning_loop_node/data_classes/image_metadata.py

@@ -46,6 +46,8 @@ class ImageMetadata():
         'description': 'Creation date of the image'})
     source: Optional[str] = field(default=None, metadata={
         'description': 'Source of the image'})
+    state: str | None = field(default=None, metadata={
+        'description': 'State the image should enter the loop in (e.g. "trash"); None uses the loop default'})
 
     def __len__(self):
         return len(self.box_detections) + len(self.point_detections) + len(self.segmentation_detections) + len(self.classification_detections)

learning_loop_node/detector/categories.py

@@ -0,0 +1,27 @@
+"""Resolving the class indices or names a model emits against ``ModelInformation.categories``."""
+
+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.
+
+    :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}')

learning_loop_node/detector/detector_node.py

@@ -18,7 +18,6 @@ from socketio import AsyncClient
 
 from ..data_classes import (
     AboutResponse,
-    Category,
     Context,
     DetectorStatus,
     ImageMetadata,
@@ -33,6 +32,7 @@ from ..globals import GLOBALS
 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
@@ -339,7 +339,9 @@ class DetectorNode(Node):
                 - bytes: bytes of the ndarray (retrieved via `ndarray.tobytes(order='C')`)
                 - dtype: data type of the ndarray as string (e.g. `uint8`, `float32`, etc.)
                 - shape: shape of the ndarray as tuple of ints (e.g. `(480, 640, 3)`)
-            - metadata: The metadata for the image (optional)
+            - metadata: The metadata for the image (optional). Besides detections and annotations it
+              may carry `tags`, `source`, `created` and `state`. A `state` of e.g. "trash" makes the
+              loop file the image into that state instead of its default one.
             - upload_priority: Whether to upload with priority (optional)
             """
             self.log.debug('Processing upload via socketio.')
@@ -707,26 +709,28 @@ class DetectorNode(Node):
             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 costs that one detection its id, not 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):

learning_loop_node/detector/geometry.py

@@ -0,0 +1,39 @@
+"""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`
+takes and produces.
+"""
+
+
+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.
+
+    :return: The clipped ``(x1, y1, width, height)``; 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_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

learning_loop_node/detector/postprocess.py

@@ -0,0 +1,219 @@
+"""Model-agnostic detection postprocessing."""
+
+import logging
+from collections.abc import Sequence
+from dataclasses import dataclass
+
+import numpy as np
+
+from ..data_classes import (
+    BoxDetection,
+    Detections,
+    ImageMetadata,
+    ModelInformation,
+    PointDetection,
+)
+from ..enums import CategoryType
+from .categories import category_by_index
+from .geometry import clip_box, clip_point
+
+logger = logging.getLogger(__name__)
+
+MIN_BOX_SIZE: int = 2
+
+
+@dataclass(kw_only=True, slots=True, frozen=True)
+class Prediction:
+    """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
+    width: float
+    height: float
+    category_index: int
+    confidence: float
+
+
+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[Prediction]:
+    """Filter by confidence, run NMS, return what survives."""
+    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)
+
+    return predictions_from_xyxy(labels=classes, boxes=boxes, scores=scores)
+
+
+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 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 predictions_from_xyxy(
+    *,
+    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."""
+    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(
+    predictions: list[Prediction],
+    model_information: ModelInformation,
+    im_height: int,
+    im_width: int,
+) -> ImageMetadata:
+    """Build the container a *detector* node reports."""
+    image_metadata = ImageMetadata()
+    _append_predictions(image_metadata, predictions, model_information, im_height, im_width)
+    return image_metadata
+
+
+def to_detections(
+    predictions: list[Prediction],
+    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_predictions(result, predictions, model_information, im_height, im_width)
+    return result
+
+
+def _append_predictions(
+    target: ImageMetadata | Detections,
+    predictions: list[Prediction],
+    model_information: ModelInformation,
+    im_height: int,
+    im_width: int,
+) -> None:
+    """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=prediction.x,
+                y1=prediction.y,
+                width=prediction.width,
+                height=prediction.height,
+                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=prediction.confidence,
+                )
+            )
+        elif category.type == CategoryType.Point:
+            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(
+                    category_name=category.name,
+                    x=cx,
+                    y=cy,
+                    category_id=category.id,
+                    model_name=model_information.version,
+                    confidence=prediction.confidence,
+                )
+            )
+        else:
+            logger.warning('Unsupported category type %s for category %s', category.type, category.name)
+
+    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)

learning_loop_node/helpers/entrypoint.py

@@ -0,0 +1,81 @@
+"""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``
+and ``NODE_PORT``: the bare ``HOST`` is the address of the loop.
+"""
+
+import logging
+import os
+from argparse import Action, Namespace
+
+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.
+
+    :param legacy_env_prefix: A prefix an earlier version of this node required, e.g.
+        ``'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',
+                        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')
+    logger.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)
+            if not name or name in os.environ:
+                continue
+            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`` is
+        ``<PREFIX>HOST``, not ``<PREFIX>NODE_HOST``.
+        """
+        candidates = [self.legacy_env_prefix + name, self.legacy_env_prefix + action.dest.upper()]
+        return list(dict.fromkeys(candidates))

learning_loop_node/helpers/environment_reader.py

@@ -1,26 +1,32 @@
 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:
+    """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))
 
     # 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}')
 
     # Possible error: multiple values are not None and not equal
     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}')
+        logger.warning('different environment variables set for %s: %s - using %s',
+                       possible_names, values, values[0])
 
     return values[0]
 

learning_loop_node/trainer/batch_size.py

@@ -0,0 +1,81 @@
+"""Choosing a batch size by probing, rather than configuring one.
+
+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.
+https://github.com/Lightning-AI/pytorch-lightning
+"""
+
+import logging
+from collections.abc import Callable
+
+from .exceptions import InsufficientMemoryError
+
+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."""
+
+MIN_TRAIN_STEPS_PER_EPOCH = 8
+"""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 InsufficientMemoryError: If not even a batch size of 1 fits.
+    """
+    limit = smaller_pot(limit)
+    if not fits(1):
+        raise InsufficientMemoryError('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.
+
+    cuDNN and cuBLAS workspace failures raise a plain ``RuntimeError``, so the message has to be
+    matched too.
+    """
+    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'))

learning_loop_node/trainer/cuda.py

@@ -0,0 +1,137 @@
+"""GPU memory budgeting and batch-size probing for trainers that train in-process.
+
+``--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
+
+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, smaller_pot
+
+logger = logging.getLogger(__name__)
+
+SAFETY_MARGIN = 0.05
+"""Share of the budget held back while probing, against allocator fragmentation later on."""
+
+
+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 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, 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.
+    """
+    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)
+    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
+
+
+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.
+
+    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).
+    """
+    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 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()

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.
-    '''
+    """

learning_loop_node/trainer/metrics.py

@@ -0,0 +1,28 @@
+"""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."""
+    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}}

learning_loop_node/trainer/subprocess.py

@@ -0,0 +1,91 @@
+"""Run a blocking, CPU-bound generator in its own process without blocking the event loop.
+
+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
+
+import asyncio
+import logging
+import multiprocessing
+import queue
+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')
+
+
+@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())
+
+
+async def _iterator_cpu_bound_inner(
+    it: Callable[P, Iterator[T]],
+    *args: P.args,
+    **kwargs: P.kwargs,
+) -> AsyncGenerator[T, None]:
+    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',
+    )
+
+    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()
+
+
+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:
+        logger.exception('iterator_cpu_bound child process failed')
+        state_queue.put(e)
+
+    state_queue.put(IteratorDone())
+
+
+class IteratorDone:
+    pass

learning_loop_node/trainer/trainer_logic.py

@@ -161,27 +161,38 @@ class TrainerLogic(TrainerLogicGeneric):
 
     @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)."""
 
-    @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]:
-        '''Called to run detections on a list of images.'''
+        """Called to run detections on a list of images."""

{learning_loop_node-0.22.0.dist-info → learning_loop_node-0.23.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: learning_loop_node
-Version: 0.22.0
+Version: 0.23.0
 Summary: Python Library for Nodes which connect to the Zauberzeug Learning Loop
 Author-email: Zauberzeug GmbH <info@zauberzeug.com>
 License: MIT
@@ -13,6 +13,7 @@ Requires-Dist: requests<3.0.0,>=2.33.0
 Requires-Dist: python-socketio<6.0.0,>=5.16.2
 Requires-Dist: aiofiles>=0.7.0
 Requires-Dist: python-multipart>=0.0.31
+Requires-Dist: configargparse>=1.7.1
 Requires-Dist: psutil<8.0.0,>=5.9.0
 Requires-Dist: numpy<3.0,>=2.0
 Requires-Dist: Pillow<13.0.0,>=12.3.0
@@ -68,6 +69,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
@@ -132,7 +137,10 @@ The detector also has a **SocketIO** upload endpoint that can be used to upload
 - `metadata`: a dictionary representing the image metadata. If metadata contains detections and/or annotations, UUIDs for the classes are automatically determined based on the category names. Metadata should follow the schema of the `ImageMetadata` data class.
 - `upload_priority`: Optional boolean flag to prioritize the upload (defaults to False)
 
-The endpoint returns None if the upload was successful and an error message otherwise.
+The endpoint returns `{'status': 'OK'}` if the upload was successful and `{'error': '<message>'}` otherwise.
+
+Besides detections and annotations, the metadata may carry `tags`, `source`, `created` and `state`.
+`state` names the state the image should enter the loop in (e.g. `trash`); when it is omitted the loop applies its own default (`inbox`).
 
 For both ways to upload an image, the tag `picked_by_system` is automatically added to the image metadata.
 

{learning_loop_node-0.22.0.dist-info → learning_loop_node-0.23.0.dist-info}/RECORD

@@ -13,14 +13,17 @@ learning_loop_node/data_classes/annotation_data.py,sha256=NfMlTv2_5AfVY_JDM4tbjE
 learning_loop_node/data_classes/annotations.py,sha256=ha7uuWmZqW-LU2vYrkvfnaVpvf5C2xdm826IVVBu0w0,1848
 learning_loop_node/data_classes/detections.py,sha256=hGSKc1elk4Drp4XUjQ3F9HGeRpyLvA0q064b-S9lIug,5335
 learning_loop_node/data_classes/general.py,sha256=y_seXgqCWGoUCFzwZKt8eG8rAlANIG8VhzVVivZLklU,7334
-learning_loop_node/data_classes/image_metadata.py,sha256=4FFs89iZAfr2Jh-40QqZ_GGmJ8hl_W61eHcCqHflJLE,2293
+learning_loop_node/data_classes/image_metadata.py,sha256=QPLJ4gcVj33xaaCFF1n5e8G6higOJ6iIKy3YhGZMfkw,2459
 learning_loop_node/data_classes/socket_response.py,sha256=tIdt-oYf6ULoJIDYQCecNM9OtWR6_wJ9tL0Ksu83Vko,655
 learning_loop_node/data_classes/training.py,sha256=TybwcCDf_NUaDUaOj30lPm-7Z3Qk9XFRibEX5qIv96Y,5737
 learning_loop_node/detector/__init__.py,sha256=wroy0W1X5MX16C-peA-UHYrcdXet7GhleIlC_Pqa6hI,128
+learning_loop_node/detector/categories.py,sha256=qgquuop3geKXXsNBurO0eC2b64h_ojL37RBYMv0ZzOg,1147
 learning_loop_node/detector/detector_logic.py,sha256=XJaHNiw5PaHoCWnJ-qsxdUELLevvztYx54TQ8KcHl18,1479
-learning_loop_node/detector/detector_node.py,sha256=As8QW2LaR81hsZePONAuf0RBNfoDEBEF18tIRM_-evs,35834
+learning_loop_node/detector/detector_node.py,sha256=GC5CaGcW_qVxyolDEz9U5W7y0ejZGEu9FZ__ckLxd7g,35728
 learning_loop_node/detector/exceptions.py,sha256=C6KbNPlSbtfgDrZx2Hbhm7Suk9jVoR3fMRCO0CkrMsQ,196
+learning_loop_node/detector/geometry.py,sha256=bZCF9xuK7qy6Ipv5aQYrTlwSmjaiUz5VEUYnc1zCe-w,1101
 learning_loop_node/detector/outbox.py,sha256=6F0MjFM-zNX-xWP_YGbmU0UYuyb0odsX_OcAxMPSzS0,12288
+learning_loop_node/detector/postprocess.py,sha256=fPBbgE6Obkej709zoFPGqA2gnbJNbpbMuZogilKy4pI,7274
 learning_loop_node/detector/inbox_filter/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 learning_loop_node/detector/inbox_filter/cam_observation_history.py,sha256=slPGm87TbBCfDUdpRTAM5yOz1An8Br5X_2HFJOGSrFg,3540
 learning_loop_node/detector/inbox_filter/relevance_filter.py,sha256=81XQV6HST2hSH-In7VF8B4BfubgUdRHTLcc16VLjUxc,1518
@@ -39,23 +42,28 @@ learning_loop_node/enums/general.py,sha256=OSpcIpdFjuRByON2xBoiIXl1SgDREIqkwnaxk
 learning_loop_node/enums/trainer.py,sha256=VaD63guLO4aKgVfXT0EryPlXKQGegSET3Cp4R0uC_Ew,746
 learning_loop_node/helpers/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 learning_loop_node/helpers/background_tasks.py,sha256=sbbS8JO1ls0ZCbNW68KOT9QLOFyFoJRXwwAq78iOmoA,3256
-learning_loop_node/helpers/environment_reader.py,sha256=6DxDJecLHxiGczByhyVa_JssAwwft7vuNCGaEzoSY2I,1662
+learning_loop_node/helpers/entrypoint.py,sha256=EtRMXZKzNapwbIylf4hg8Obw-2EydrLGKUjbIn5H-94,3591
+learning_loop_node/helpers/environment_reader.py,sha256=zsoFysxNxBCOJzUwX4Pbm3ftYhxx5zsbOHQfD9VcfvA,1937
 learning_loop_node/helpers/gdrive_downloader.py,sha256=zeYJciTAJVRpu_eFjwgYLCpIa6hU1d71anqEBb564Rk,1145
 learning_loop_node/helpers/log_conf.py,sha256=hqVAa_9NnYEU6N0dcOKmph82p7MpgKqeF_eomTLYzWY,961
 learning_loop_node/helpers/misc.py,sha256=qdN18Ukbw6eR6sFL4BxNq0EZzg2tVm2Cm8cTbFyy8QU,8931
 learning_loop_node/helpers/run.py,sha256=_uox-j3_K_bL3yCAwy3JYSOiIxrnhzVxyxWpCe8_J9U,876
 learning_loop_node/trainer/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+learning_loop_node/trainer/batch_size.py,sha256=Ma_nH3PHAk2gzqdkNXstNFqwXUUTt7pVe4aqWRrkHX4,2914
+learning_loop_node/trainer/cuda.py,sha256=LkXJaJ404S64XGRTMqqULB9e6g85dJfjLIAn12lnpTY,5968
 learning_loop_node/trainer/downloader.py,sha256=AIyEKM4XYtd6VgmXrP3VayV9DpJzdURK1Brx81ePNSM,1470
-learning_loop_node/trainer/exceptions.py,sha256=vbuoE6kssLQuA8zd3LiDHmZglP6E2IJJwEi5AZtWXxY,420
+learning_loop_node/trainer/exceptions.py,sha256=IRT7qvn7-RjXaqL8mF21ySm7JWhQtqm1Cl0Gi3DChVY,540
 learning_loop_node/trainer/executor.py,sha256=XQ1TwgOF2pQv7bgzSaljDWG4I1ySjYXL8r9c_p9MX1Q,3947
 learning_loop_node/trainer/io_helpers.py,sha256=zWl0-HLP8zxOG3ybm76vXgEozULj_1QRl95jkrF126k,7255
+learning_loop_node/trainer/metrics.py,sha256=H1jQh5fHN3-cTdMTuH8mF_uODV0T9DBdqcPUD-ZqFAI,1439
+learning_loop_node/trainer/subprocess.py,sha256=WxuTly0zhpClEjTd9wqoI_-7YTNMaiZpRHxu00vttKw,2647
 learning_loop_node/trainer/test_executor.py,sha256=6BVGDN_6f5GEMMEvDLSG1yzMybSvgXaP5uYpSfsVPP0,2224
-learning_loop_node/trainer/trainer_logic.py,sha256=BgkeDEVaRVuRETYMt2rKe5cspExTgnO7LXNVYVbIcsw,8679
+learning_loop_node/trainer/trainer_logic.py,sha256=Yaw1GNTC05rVLQczY6vKUgkymBYfm3_0XaR5C0q-CP4,9159
 learning_loop_node/trainer/trainer_logic_generic.py,sha256=IRP68ZGpnOX4OS03OwB7HmTUvk6nbBRmCjMCVupGE_Q,27196
 learning_loop_node/trainer/trainer_node.py,sha256=ZnydfHJ2ehAhb821Xke_09qthCE3vDf7iqeYUrliSSo,4938
 learning_loop_node/trainer/rest/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 learning_loop_node/trainer/rest/backdoor_controls.py,sha256=ZnK8ypY5r_q0-YZbtaOxhQThzuZvMsQHM5gJGESd_dE,5131
-learning_loop_node-0.22.0.dist-info/METADATA,sha256=AuMYnVs2JgNTb3LNoMJmG0XdjhU2dUpOH_y7L-VpYMk,17030
-learning_loop_node-0.22.0.dist-info/WHEEL,sha256=YVMoNqKzERt-wjUZwJ33xBGAwnFl-4cqbYkTtWa4itE,91
-learning_loop_node-0.22.0.dist-info/top_level.txt,sha256=nS5qiFc4stVF_VXVCVJxZnrqVOv3IGX5-_E3sLby6OE,19
-learning_loop_node-0.22.0.dist-info/RECORD,,
+learning_loop_node-0.23.0.dist-info/METADATA,sha256=L8ASHsKUg7Z_JGji4tqDJEbsCf44tZoKiPsZ4YeU6oo,17607
+learning_loop_node-0.23.0.dist-info/WHEEL,sha256=YVMoNqKzERt-wjUZwJ33xBGAwnFl-4cqbYkTtWa4itE,91
+learning_loop_node-0.23.0.dist-info/top_level.txt,sha256=nS5qiFc4stVF_VXVCVJxZnrqVOv3IGX5-_E3sLby6OE,19
+learning_loop_node-0.23.0.dist-info/RECORD,,