learning-loop-node 0.22.0__tar.gz → 0.23.0__tar.gz

{learning_loop_node-0.22.0 → learning_loop_node-0.23.0}/PKG-INFO

@@ -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 → learning_loop_node-0.23.0}/README.md

@@ -36,6 +36,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
@@ -100,7 +104,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 → learning_loop_node-0.23.0}/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-0.23.0/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-0.22.0 → learning_loop_node-0.23.0}/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-0.23.0/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-0.23.0/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-0.23.0/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-0.22.0 → learning_loop_node-0.23.0}/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-0.23.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'))