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

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: learning_loop_node
-Version: 0.23.0
+Version: 0.25.0
 Summary: Python Library for Nodes which connect to the Zauberzeug Learning Loop
 Author-email: Zauberzeug GmbH <info@zauberzeug.com>
 License: MIT
@@ -60,11 +60,14 @@ You can configure connection to our Learning Loop by specifying the following en
 | LOOP_PROJECT             | PROJECT      | Project ID                                                   | Detector                  | -            |
 | MIN_UNCERTAIN_THRESHOLD  | -            | smallest confidence (float) at which auto-upload will happen | Detector (opt.)           | 0.3          |
 | MAX_UNCERTAIN_THRESHOLD  | -            | largest confidence (float) at which auto-upload will happen  | Detector (opt.)           | 0.6          |
+| MAX_OUTBOX_SIZE          | -            | max. number of non-priority images queued for upload         | Detector (opt.)           | 1000         |
 | EXCLUSIVE_MODEL_BUILD    | -            | Reject detections during update to save VRAM (set to 1)      | Detector (opt.)           | 0            |
 | INFERENCE_BATCH_SIZE     | -            | Batch size of trainer when calculating detections            | Trainer (opt.)            | 10           |
+| VRAM_LIMIT_GB            | -            | GPU memory (GB) a training may use; the batch size is probed against it (`--vram-limit-gb`, trainers built with `node_parser(vram_limit=True)`) | Trainer (opt.) | 0 (whole card) |
 | RESTART_AFTER_TRAINING   | -            | Restart the trainer after training (set to 1)                | Trainer (opt.)            | 0            |
 | KEEP_OLD_TRAININGS       | -            | Do not delete old trainings (set to 1)                       | Trainer (opt.)            | 0            |
 | TRAINER_IDLE_TIMEOUT_SEC | -            | Automatically shutdown trainer after timeout (in seconds)    | Trainer (opt.)            | 0 (disabled) |
+| TRAINING_SLOT_LOCK       | -            | Lock file shared by the trainers on one GPU (see Trainer Node) | Trainer (opt.)          | - (disabled) |
 | USE_BACKDOOR_CONTROLS    | -            | Always enable backdoor controls (set to 1)                   | Trainer / Detector (opt.) | 0            |
 
 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.
@@ -73,6 +76,16 @@ Where a name has an alias, either spelling works. If both are set to **different
 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`.
 
+#### Debug logging
+
+Every node type (Detector, Trainer, Annotator, Converter) has a REST endpoint to switch its own log level between `INFO` and `DEBUG` at runtime. Send the plain text `on` or `off` with a PUT request:
+
+- Turn debug logging on: `curl -X PUT -d "on" http://localhost/debug_logging`
+- Turn debug logging off: `curl -X PUT -d "off" http://localhost/debug_logging`
+
+The response repeats the new state. Any other body returns status code 400.
+Only the node's own logger is affected; third-party loggers such as `httpx` keep their level. The setting is not persistent, a restart brings the log level back to `INFO`.
+
 #### Testing
 
 We use github actions for CI. Tests can also be executed locally by running
@@ -238,6 +251,20 @@ Trainers fetch the images and anntoations from the Learning Loop to train new mo
 
 - if the command line tool "jpeginfo" is installed, the downloader will drop corrupted images automatically
 
+### Several trainers on one GPU
+
+Trainers that share a GPU (say, one per model architecture on the same machine) coordinate through
+a lock file: mount one host directory into each of them and point `TRAINING_SLOT_LOCK` at a file
+inside it, e.g. `/slot/training.lock`. A trainer takes the lock when it accepts a training and holds
+it until the training is cleaned up, including detection and upload. While a sibling holds it, the
+others report the state `blocked` instead of `idle`, so the Learning Loop neither offers them for a
+training nor dispatches queued jobs to them; a trainer that is sent `begin_training` regardless
+refuses it with `{"success": false, "error_msg": ...}` in the acknowledgement and sets nothing up. A
+training resumed after a restart cannot be refused — it waits for the lock instead, reporting
+`waiting_for_slot` meanwhile. The lock is an `flock`, so the kernel releases it whenever the holding
+process ends, however it ends — there is nothing to clean up after a crash. Mount the directory,
+not the file, and never replace the file: the lock belongs to its inode.
+
 ## Converter Node
 
 A Conveter Node converts models from one format into another.

{learning_loop_node-0.23.0 → learning_loop_node-0.25.0}/README.md

@@ -27,11 +27,14 @@ You can configure connection to our Learning Loop by specifying the following en
 | LOOP_PROJECT             | PROJECT      | Project ID                                                   | Detector                  | -            |
 | MIN_UNCERTAIN_THRESHOLD  | -            | smallest confidence (float) at which auto-upload will happen | Detector (opt.)           | 0.3          |
 | MAX_UNCERTAIN_THRESHOLD  | -            | largest confidence (float) at which auto-upload will happen  | Detector (opt.)           | 0.6          |
+| MAX_OUTBOX_SIZE          | -            | max. number of non-priority images queued for upload         | Detector (opt.)           | 1000         |
 | EXCLUSIVE_MODEL_BUILD    | -            | Reject detections during update to save VRAM (set to 1)      | Detector (opt.)           | 0            |
 | INFERENCE_BATCH_SIZE     | -            | Batch size of trainer when calculating detections            | Trainer (opt.)            | 10           |
+| VRAM_LIMIT_GB            | -            | GPU memory (GB) a training may use; the batch size is probed against it (`--vram-limit-gb`, trainers built with `node_parser(vram_limit=True)`) | Trainer (opt.) | 0 (whole card) |
 | RESTART_AFTER_TRAINING   | -            | Restart the trainer after training (set to 1)                | Trainer (opt.)            | 0            |
 | KEEP_OLD_TRAININGS       | -            | Do not delete old trainings (set to 1)                       | Trainer (opt.)            | 0            |
 | TRAINER_IDLE_TIMEOUT_SEC | -            | Automatically shutdown trainer after timeout (in seconds)    | Trainer (opt.)            | 0 (disabled) |
+| TRAINING_SLOT_LOCK       | -            | Lock file shared by the trainers on one GPU (see Trainer Node) | Trainer (opt.)          | - (disabled) |
 | USE_BACKDOOR_CONTROLS    | -            | Always enable backdoor controls (set to 1)                   | Trainer / Detector (opt.) | 0            |
 
 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.
@@ -40,6 +43,16 @@ Where a name has an alias, either spelling works. If both are set to **different
 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`.
 
+#### Debug logging
+
+Every node type (Detector, Trainer, Annotator, Converter) has a REST endpoint to switch its own log level between `INFO` and `DEBUG` at runtime. Send the plain text `on` or `off` with a PUT request:
+
+- Turn debug logging on: `curl -X PUT -d "on" http://localhost/debug_logging`
+- Turn debug logging off: `curl -X PUT -d "off" http://localhost/debug_logging`
+
+The response repeats the new state. Any other body returns status code 400.
+Only the node's own logger is affected; third-party loggers such as `httpx` keep their level. The setting is not persistent, a restart brings the log level back to `INFO`.
+
 #### Testing
 
 We use github actions for CI. Tests can also be executed locally by running
@@ -205,6 +218,20 @@ Trainers fetch the images and anntoations from the Learning Loop to train new mo
 
 - if the command line tool "jpeginfo" is installed, the downloader will drop corrupted images automatically
 
+### Several trainers on one GPU
+
+Trainers that share a GPU (say, one per model architecture on the same machine) coordinate through
+a lock file: mount one host directory into each of them and point `TRAINING_SLOT_LOCK` at a file
+inside it, e.g. `/slot/training.lock`. A trainer takes the lock when it accepts a training and holds
+it until the training is cleaned up, including detection and upload. While a sibling holds it, the
+others report the state `blocked` instead of `idle`, so the Learning Loop neither offers them for a
+training nor dispatches queued jobs to them; a trainer that is sent `begin_training` regardless
+refuses it with `{"success": false, "error_msg": ...}` in the acknowledgement and sets nothing up. A
+training resumed after a restart cannot be refused — it waits for the lock instead, reporting
+`waiting_for_slot` meanwhile. The lock is an `flock`, so the kernel releases it whenever the holding
+process ends, however it ends — there is nothing to clean up after a crash. Mount the directory,
+not the file, and never replace the file: the lock belongs to its inode.
+
 ## Converter Node
 
 A Conveter Node converts models from one format into another.

{learning_loop_node-0.23.0 → learning_loop_node-0.25.0}/learning_loop_node/detector/outbox.py

@@ -37,7 +37,7 @@ class Outbox():
     It handles upload failures by splitting the upload into two smaller batches until the problematic image is identified - and removed.
     Any image can be saved to the normal or the priority queue.
     Images in the priority queue are uploaded first.
-    The total queue length is limited to 1000 images.
+    The normal queue is limited to MAX_OUTBOX_SIZE images (default 1000); the oldest are dropped first.
     """
 
     def __init__(self) -> None:
@@ -59,7 +59,9 @@ class Outbox():
         self.log.info('Outbox initialized with target_uri: %s', self.target_uri)
 
         self.BATCH_SIZE = 20
-        self.MAX_UPLOAD_LENGTH = 1000  # only affects the `upload_folders` list
+        self.MAX_UPLOAD_LENGTH = int(os.environ.get('MAX_OUTBOX_SIZE', '1000'))  # only limits `upload_folders`
+        if self.MAX_UPLOAD_LENGTH <= 0:
+            raise ValueError(f'MAX_OUTBOX_SIZE must be positive, got {self.MAX_UPLOAD_LENGTH}')
         self.UPLOAD_INTERVAL_S = 5
         self.UPLOAD_TIMEOUT_S = 30
 

{learning_loop_node-0.23.0 → learning_loop_node-0.25.0}/learning_loop_node/enums/trainer.py

@@ -4,7 +4,9 @@ from enum import Enum
 
 class TrainerState(str, Enum):
     Idle = 'idle'
+    Blocked = 'blocked'  # idle, but unable to take a training
     Initialized = 'initialized'
+    WaitingForSlot = 'waiting_for_slot'
     Preparing = 'preparing'
     DataDownloading = 'data_downloading'
     DataDownloaded = 'data_downloaded'

{learning_loop_node-0.23.0 → learning_loop_node-0.25.0}/learning_loop_node/helpers/entrypoint.py

@@ -12,21 +12,27 @@ from argparse import Action, Namespace
 import configargparse
 import uvicorn
 
+from ..trainer.batch_size import VRAM_LIMIT_GB_FLAG, VRAM_LIMIT_GB_HELP
+
 logger = logging.getLogger(__name__)
 
 
-def node_parser(*, description: str, legacy_env_prefix: str = '') -> configargparse.ArgumentParser:
+def node_parser(*, description: str, legacy_env_prefix: str = '',
+                vram_limit: bool = False) -> 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.
+    :param vram_limit: Add :data:`~learning_loop_node.trainer.batch_size.VRAM_LIMIT_GB_FLAG`.
     """
     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')
+    if vram_limit:
+        parser.add_argument(VRAM_LIMIT_GB_FLAG, type=float, default=0, help=VRAM_LIMIT_GB_HELP)
     return parser
 
 

{learning_loop_node-0.23.0 → learning_loop_node-0.25.0}/learning_loop_node/helpers/misc.py

@@ -150,7 +150,7 @@ def ensure_socket_response(func):
             if isinstance(value, bool):
                 return asdict(SocketResponse.from_bool(value))
             if isinstance(value, SocketResponse):
-                return value
+                return asdict(value)
             if (args[0] in ['connect', 'disconnect', 'connect_error']):
                 return value
             if value is None:

learning_loop_node-0.25.0/learning_loop_node/trainer/batch_size.py

@@ -0,0 +1,126 @@
+"""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, Mapping
+from typing import Any
+
+from .exceptions import InsufficientMemoryError
+
+logger = logging.getLogger(__name__)
+
+REQUESTED_BATCH_SIZE = 'max_batch_size'
+"""The hyperparameter every node reads its `measure_batch_size` argument out of.
+
+0 or absent means the card decides. It is an input only: a trainer reports the size it settled on
+under a different key, conventionally ``batch_size``, and never writes it back here.
+"""
+
+VRAM_LIMIT_GB_FLAG = '--vram-limit-gb'
+"""The flag a trainer takes its GPU budget from; ``VRAM_LIMIT_GB`` follows from the name."""
+
+VRAM_LIMIT_GB_HELP = ('Gigabytes of GPU memory a training may use. The batch size is probed against this limit '
+                      'instead of the whole card, so a lower limit yields a smaller batch size rather than an '
+                      'out-of-memory error. Use it to share a GPU or to keep headroom against fragmentation. '
+                      "The limit is relative to the card's total memory, not to what is currently free. "
+                      '0 (default) means no limit.')
+
+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, minimum: int = 1,
+                    candidate: int = 0) -> int:
+    """Return the largest batch size that fits, never exceeding ``limit``.
+
+    Powers of two, plus ``minimum`` and ``candidate``, the two sizes outside that set this will
+    return, and only because somebody named them and they then measured.
+
+    :param fits: Runs a representative probe; ``False`` on out-of-memory.
+    :param limit: Upper bound; the doubling stops at the largest power of two within it.
+    :param minimum: Smallest size to try, as named; the doubling continues from the next power of
+        two above it. Raise it above one for a step that cannot run on a single sample at all:
+        BatchNorm over a 1x1 feature map, a validation pass that halves the batch.
+    :param candidate: An exact size, tried once the doubling has reached its ceiling, so a size
+        that was asked for is used as asked for rather than rounded down. Ignored unless it lies
+        between that ceiling and ``limit``; a bound nobody named — one derived from the dataset,
+        say — must not be passed here.
+    :raises InsufficientMemoryError: If not even ``minimum`` fits.
+    """
+    minimum = max(1, minimum)
+    bound = max(limit, minimum)
+    ceiling = max(smaller_pot(bound), minimum)
+
+    if not fits(minimum):
+        raise InsufficientMemoryError(f'batch size {minimum} does not fit in memory')
+
+    size = minimum
+    while size < ceiling and fits(smaller_pot(size) * 2):
+        size = smaller_pot(size) * 2
+
+    if size == ceiling and ceiling < candidate <= bound and fits(candidate):
+        size = candidate
+
+    return size
+
+
+def requested_batch_size(hyperparameters: Mapping[str, Any]) -> int:
+    """The bound a training asked for, read out of the hyperparameters the loop sent.
+
+    An unfilled field arrives as absent, ``None`` or ``''``, and all three mean no bound.
+
+    :raises ValueError: If the value is there but is not a number.
+    """
+    return int(hyperparameters.get(REQUESTED_BATCH_SIZE, 0) or 0)
+
+
+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, minimum: int = 1) -> int:
+    """The batch size to fall back on when there is no GPU to probe; never below ``minimum``."""
+    batch_size = max(min(smaller_pot(limit), NO_GPU_BATCH_SIZE), minimum)
+    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-0.25.0/learning_loop_node/trainer/cuda.py

@@ -0,0 +1,279 @@
+"""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 abc import ABC, abstractmethod
+from argparse import ArgumentParser
+from collections.abc import Callable
+
+import torch
+
+from .batch_size import (
+    MAX_BATCH_SIZE,
+    REQUESTED_BATCH_SIZE,
+    VRAM_LIMIT_GB_FLAG,
+    VRAM_LIMIT_GB_HELP,
+    dataset_limit,
+    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."""
+
+DEFAULT_STEPS_PER_TRIAL = 4
+"""Training steps a batch size must survive before it counts as fitting.
+
+Some trainings reach their peak only after a few steps: an optimizer that allocates its state on
+the first one, gradients that accumulate, a scaler that skips the first updates.
+"""
+
+_budget_gb: float = 0
+"""What :func:`limit_cuda_memory` capped this process to; 0 means the whole card."""
+
+
+class ProbeStep(ABC):
+    """A training cycle built for one batch-size probe and released once the probe is done.
+
+    :func:`measure_batch_size` builds it through a factory, after the safety margin is reserved,
+    and holds the only reference to it, so nothing built for the probe stays on the card for the
+    training that follows. A trial at one batch size is ``steps_per_trial`` calls of
+    :meth:`train_step` followed by one of :meth:`val_step`; a subclass implements both, the other
+    two are optional.
+    """
+
+    @abstractmethod
+    def train_step(self, batch_size: int) -> str | None:
+        """Run one training step at ``batch_size``; may return a detail to append to the log line."""
+
+    @abstractmethod
+    def val_step(self, batch_size: int) -> str | None:
+        """Run the validation a training at ``batch_size`` runs between epochs.
+
+        Validate at the batch size the training validates with, which need not be ``batch_size``
+        (YOLOv5 halves it). A training that does not validate returns ``None`` without running
+        anything.
+        """
+
+    def on_out_of_memory(self) -> None:  # noqa: B027
+        """Drop what a trial that ran out of memory left behind (an optimizer's gradients, say)."""
+
+    def release(self) -> None:  # noqa: B027
+        """Undo what building the step changed outside it, such as moving the real model off the card."""
+
+
+def measure_batch_size(build_step: Callable[[], ProbeStep], *, max_batch_size: int = 0,
+                       sample_count: int | None = None, probe: str = 'batch-size probe',
+                       minimum: int = 1, steps_per_trial: int = DEFAULT_STEPS_PER_TRIAL) -> int:
+    """Settle a training's batch size against what it asked for and what the card allows.
+
+    Every training enters here, whatever shape its hyperparameters have; :func:`probe_batch_size` is
+    for a probe with no requested size to honour, such as a detection pass bounded only by how many
+    images there are. The step is built only once there is a card to probe.
+
+    :param build_step: Builds the step the trials run; called once, after the margin is reserved.
+    :param max_batch_size: What the training asked for, as carried in the
+        :data:`~.batch_size.REQUESTED_BATCH_SIZE` hyperparameter: the largest batch it may use,
+        measured rather than trusted. A size that fits is used as asked for, whether or not it is a
+        power of two; one that does not becomes the largest power of two below it that does. 0 means
+        the card decides alone. A request below ``minimum`` is raised to it, with a warning.
+    :param sample_count: Samples in the training split, when the caller knows it. The search is
+        then bounded by :func:`~.batch_size.dataset_limit`; that bound is never used as the exact
+        candidate.
+    :param probe: Names this probe in the log, so a node running several stays readable.
+    :param minimum: Smallest size to try, taking precedence over ``max_batch_size``; see
+        :func:`~.batch_size.find_batch_size`.
+    :param steps_per_trial: Training steps a size must survive, each trial then ending in one
+        validation step.
+    :raises InsufficientMemoryError: If not even ``minimum`` fits.
+    :raises ValueError: If the training asked for a negative batch size, or ``steps_per_trial`` is
+        below one.
+    """
+    if max_batch_size < 0:
+        raise ValueError(f'{REQUESTED_BATCH_SIZE} must be >= 0, got {max_batch_size}')
+    if steps_per_trial < 1:
+        raise ValueError(f'steps_per_trial must be >= 1, got {steps_per_trial}')
+    if 0 < max_batch_size < minimum:
+        logger.warning('%s: requested %s=%d is below the trainer minimum of %d; using %d', probe,
+                       REQUESTED_BATCH_SIZE, max_batch_size, minimum, minimum)
+
+    limit = max_batch_size
+    if sample_count is not None:
+        limit = min(limit or MAX_BATCH_SIZE, dataset_limit(sample_count))
+        logger.info('%s: %d training samples allow at most %d per batch', probe, sample_count,
+                    dataset_limit(sample_count))
+
+    return _probe(build_step, probe=probe, limit=limit, steps_per_trial=steps_per_trial,
+                  candidate=max_batch_size, minimum=minimum)
+
+
+def probe_batch_size(run_batch: Callable[[int], str | None], *, probe: str = 'batch-size probe',
+                     limit: int = 0) -> int:
+    """Find the largest batch that fits for a pass with no requested size, such as a detection pass.
+
+    :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`.
+    :raises InsufficientMemoryError: If not even a batch of one fits.
+    """
+    return _probe(lambda: _PassStep(run_batch), probe=probe, limit=limit, steps_per_trial=1)
+
+
+class _PassStep(ProbeStep):
+    """A plain callable as a step; there is nothing of its own to build, validate or release."""
+
+    def __init__(self, run_batch: Callable[[int], str | None]) -> None:
+        self._run_batch = run_batch
+
+    def train_step(self, batch_size: int) -> str | None:
+        return self._run_batch(batch_size)
+
+    def val_step(self, batch_size: int) -> None:
+        return None
+
+
+def _probe(build_step: Callable[[], ProbeStep], *, probe: str, limit: int, steps_per_trial: int,
+           candidate: int = 0, minimum: int = 1) -> int:
+    """Everything of a probe but the step: the GPU check, the margin, the search, the release.
+
+    The margin is a share of the budget :func:`limit_cuda_memory` set in this process.
+    """
+    bound = max(limit or MAX_BATCH_SIZE, minimum)
+
+    if not torch.cuda.is_available():
+        return no_gpu_batch_size(bound, probe, minimum)
+
+    margin = reserve_margin(probe=probe)
+    try:
+        chosen = _search(build_step, probe=probe, bound=bound, steps_per_trial=steps_per_trial,
+                         minimum=minimum, candidate=candidate)
+    finally:
+        del margin
+        free_cuda_memory()
+    logger.info('%s: selected batch size %d (upper bound %d)', probe, chosen, bound)
+    return chosen
+
+
+def _search(build_step: Callable[[], ProbeStep], *, probe: str, bound: int, steps_per_trial: int,
+            minimum: int, candidate: int) -> int:
+    """Build the step, search with it and release it; once this returns, nothing references it."""
+    step = build_step()
+    try:
+        fits = measured_fits(_trial(step, steps_per_trial), probe=probe, on_out_of_memory=step.on_out_of_memory)
+        return find_batch_size(fits, limit=bound, minimum=minimum, candidate=candidate)
+    finally:
+        step.release()
+
+
+def _trial(step: ProbeStep, steps_per_trial: int) -> Callable[[int], str | None]:
+    """One trial of ``step``: its training steps, then its validation."""
+    def run_trial(batch_size: int) -> str | None:
+        train_detail = None
+        for _ in range(steps_per_trial):
+            train_detail = step.train_step(batch_size)
+        val_detail = step.val_step(batch_size)
+        return '; '.join(detail for detail in (train_detail, val_detail) if detail) or None
+
+    return run_trial
+
+
+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(*, 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.
+    """
+    margin_bytes = int(usable_memory_bytes() * 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() -> int:
+    """How much GPU memory this process may allocate, honouring :func:`limit_cuda_memory`."""
+    if _budget_gb > 0:
+        return int(_budget_gb * 1024**3)
+    return torch.cuda.get_device_properties(None).total_memory
+
+
+def add_vram_limit_argument(parser: ArgumentParser) -> None:
+    """Give a spawned training script the same GPU budget flag its node has.
+
+    The spawned process still has to call :func:`limit_cuda_memory` with it.
+    """
+    parser.add_argument(VRAM_LIMIT_GB_FLAG, type=float, default=0, help=VRAM_LIMIT_GB_HELP)
+
+
+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. A probe in this process then measures against the same budget.
+
+    :param vram_limit_gb: 0 or less means no cap.
+    """
+    global _budget_gb  # pylint: disable=global-statement
+    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)
+    _budget_gb = vram_limit_gb
+    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-0.23.0 → learning_loop_node-0.25.0}/learning_loop_node/trainer/exceptions.py

@@ -5,6 +5,18 @@ class CriticalError(Exception):
     """
 
 
+class TrainingRefusedError(Exception):
+    """Raised when a trainer cannot take on a training it was asked to begin; nothing was set up."""
+
+
+class TrainingSlotTakenError(TrainingRefusedError):
+    """Raised when a sibling node holds the training slot."""
+
+    def __init__(self, holder: str) -> None:
+        super().__init__(f'the training slot is held by {holder}')
+        self.holder = holder
+
+
 class InsufficientMemoryError(RuntimeError):
     """Raised when not even the smallest unit of work fits in memory."""