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

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/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/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/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/trainer/batch_size.py

@@ -10,12 +10,29 @@ https://github.com/Lightning-AI/pytorch-lightning
 """
 
 import logging
-from collections.abc import Callable
+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."""
 
@@ -26,23 +43,51 @@ 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``.
+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.
-    :raises InsufficientMemoryError: If not even a batch size of 1 fits.
+    :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.
     """
-    limit = smaller_pot(limit)
-    if not fits(1):
-        raise InsufficientMemoryError('batch size 1 does not fit in memory')
+    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 = 1
-    while size < limit and fits(size * 2):
-        size *= 2
+    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:
@@ -50,9 +95,9 @@ def dataset_limit(sample_count: int) -> int:
     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)
+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
 

learning_loop_node/trainer/cuda.py

@@ -11,47 +11,185 @@ 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, find_batch_size, is_out_of_memory, no_gpu_batch_size, smaller_pot
+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, 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.
+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, rounded down to a power of two; 0 means
+    :param limit: Caps the search; 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.
+    :raises InsufficientMemoryError: If not even a batch of one fits.
     """
-    limit = smaller_pot(limit or MAX_BATCH_SIZE)
+    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(limit, probe)
+        return no_gpu_batch_size(bound, probe, minimum)
 
-    margin = reserve_margin(vram_limit_gb, probe=probe)
+    margin = reserve_margin(probe=probe)
     try:
-        chosen = find_batch_size(measured_fits(run_batch, probe=probe), limit=limit)
+        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, limit)
+    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.
@@ -84,38 +222,41 @@ def measured_fits(run_batch: Callable[[int], str | None], *, probe: str,
     return fits
 
 
-def reserve_margin(vram_limit_gb: float, *, probe: str) -> torch.Tensor:
+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.
-
-    :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)
+    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(vram_limit_gb: float) -> int:
-    """How much GPU memory this process may allocate, honouring :func:`limit_cuda_memory`.
+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
+
 
-    :param vram_limit_gb: 0 or less means the whole card.
+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.
     """
-    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))
+    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.
+    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
 
@@ -129,6 +270,7 @@ def limit_cuda_memory(vram_limit_gb: float) -> None:
         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)
 
 

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

learning_loop_node/trainer/trainer_logic_generic.py

@@ -26,8 +26,9 @@ from ..helpers.misc import (
     is_valid_uuid4,
 )
 from .downloader import TrainingsDownloader
-from .exceptions import CriticalError, NodeNeedsRestartError
+from .exceptions import CriticalError, NodeNeedsRestartError, TrainingRefusedError
 from .io_helpers import ActiveTrainingIO, EnvironmentVars, LastTrainingIO
+from .training_slot import TrainingSlot, training_slot_from_env
 
 if TYPE_CHECKING:
     from .trainer_node import TrainerNode
@@ -37,7 +38,7 @@ logger = logging.getLogger('learning_loop_node.trainer_logic_generic')
 
 class TrainerLogicGeneric(ABC):
 
-    def __init__(self, model_format: str) -> None:
+    def __init__(self, model_format: str, training_slot: TrainingSlot | None = None) -> None:
 
         # NOTE: model_format is used in the file path for the model on the server:
         # It acts as a key for list of files (cf. _get_latest_model_files)
@@ -54,6 +55,7 @@ class TrainerLogicGeneric(ABC):
         self._training: Optional[Training] = None
         self._active_training_io: Optional[ActiveTrainingIO] = None
         self._environment_vars = EnvironmentVars()
+        self.training_slot = training_slot or training_slot_from_env()
 
     # ---------------------------------------- PROPERTIES TO AVOID CHECKING FOR NONE ----------------------------------------
 
@@ -98,9 +100,17 @@ class TrainerLogicGeneric(ABC):
 
     @property
     def state(self) -> str:
-        """Returns the current state of the training. Used solely by the node in send_status().
+        """Returns the current state of the training.
+
+        Without a training the state is `idle`, or `blocked` while a sibling node holds the training slot.
+        A training resumed after a restart reports `waiting_for_slot` until it gets the slot: it has a
+        training, so it is never `blocked`.
         """
-        if (not self.training_active) or (self.training.training_state is None):
+        if not self.training_active:
+            return TrainerState.Blocked.value if self.training_slot.taken_by_sibling else TrainerState.Idle.value
+        if not self.training_slot.held:
+            return TrainerState.WaitingForSlot.value
+        if self.training.training_state is None:
             return TrainerState.Idle.value
         return self.training.training_state
 
@@ -213,8 +223,21 @@ class TrainerLogicGeneric(ABC):
 
     async def begin_training(self, organization: str, project: str, training_config: Dict) -> None:
         """Called on `begin_training` event from the Learning Loop.
+
+        The training slot is taken before anything is set up, so a refused training leaves the
+        previous training folders in place.
+
+        :raise TrainingRefusedError: if a training is active or a sibling node holds the training slot.
         """
-        self._init_new_training(Context(organization=organization, project=project), training_config)
+        if self.training_active:
+            raise TrainingRefusedError(f'training {self.training.id} is still active')
+        context = Context(organization=organization, project=project)
+        self.training_slot.claim(self._slot_claim(context, training_config.get('training_number')))
+        try:
+            self._init_new_training(context, training_config)
+        except Exception:
+            self.training_slot.release()
+            raise
         self._begin_training_task()
 
     def _begin_training_task(self) -> None:
@@ -242,6 +265,7 @@ class TrainerLogicGeneric(ABC):
         """
         self.errors.reset_all()
         try:
+            await self.training_slot.wait_for(self._slot_claim(self.training.context, self.training.training_number))
             await self._training_loop()
         except asyncio.CancelledError:
             if not self.shutdown_event.is_set():
@@ -254,6 +278,12 @@ class TrainerLogicGeneric(ABC):
                 logger.info('CancelledError in _run - shutting down')
         except Exception:
             logger.exception('(Ignored) exception in trainer_logic._run:')
+        finally:
+            self._training = None
+            self.training_slot.release()
+
+    def _slot_claim(self, context: Context, training_number: Any) -> str:
+        return f'node {self.node.uuid}, training #{training_number} of {context.organization}/{context.project}'
 
     # ---------------------------------------- TRAINING STATES ----------------------------------------
 

learning_loop_node/trainer/trainer_node.py

@@ -7,7 +7,10 @@ from typing import Dict, Optional
 from fastapi.encoders import jsonable_encoder
 from socketio import AsyncClient, exceptions
 
+from ..data_classes import SocketResponse
+from ..enums import TrainerState
 from ..node import Node
+from .exceptions import TrainingRefusedError
 from .io_helpers import LastTrainingIO
 from .rest import backdoor_controls
 from .trainer_logic_generic import TrainerLogicGeneric
@@ -64,7 +67,11 @@ class TrainerNode(Node):
         @sio_client.event
         async def begin_training(organization: str, project: str, details: Dict):
             self.log.info('received begin_training from server')
-            await self.trainer_logic.begin_training(organization, project, details)
+            try:
+                await self.trainer_logic.begin_training(organization, project, details)
+            except TrainingRefusedError as e:
+                self.log.warning('refused begin_training: %s', e)
+                return SocketResponse(success=False, error_msg=str(e), payload='refused')
             return True
 
         @sio_client.event
@@ -102,7 +109,7 @@ class TrainerNode(Node):
         if not self._idle_timeout:
             return
 
-        if self.trainer_logic.state == 'idle':
+        if self.trainer_logic.state in (TrainerState.Idle, TrainerState.Blocked):
             if self._first_idle_time is None:
                 self._first_idle_time = time.time()
             idle_time = time.time() - self._first_idle_time

learning_loop_node/trainer/training_slot.py

@@ -0,0 +1,207 @@
+"""The right to use the GPU, shared between the trainer nodes on one machine.
+
+A node holds the slot from the moment it accepts a training until the training is cleared.
+"""
+import asyncio
+import fcntl
+import logging
+import os
+import socket
+import weakref
+from abc import ABC, abstractmethod
+from pathlib import Path
+
+from .exceptions import TrainingSlotTakenError
+
+ENV_VAR = 'TRAINING_SLOT_LOCK'
+_CLAIM_ATTEMPTS = 3
+
+logger = logging.getLogger('learning_loop_node.training_slot')
+
+
+class TrainingSlot(ABC):
+    """The slot as a trainer uses it; an adapter implements `_acquire`, `_sibling_holder`, `release` and `held`."""
+
+    def claim(self, training: str) -> None:
+        """Take the slot for `training` without waiting; a no-op if this slot already holds it.
+
+        :raise TrainingSlotTakenError: if a sibling node holds the slot.
+        """
+        holder = self._describe(training)
+        for _ in range(_CLAIM_ATTEMPTS - 1):
+            try:
+                self._acquire(holder)
+                return
+            except TrainingSlotTakenError:
+                if self._sibling_holder() is not None:
+                    raise
+        self._acquire(holder)
+
+    async def wait_for(self, training: str) -> None:
+        """Take the slot for `training`, waiting as long as a sibling node holds it."""
+        try:
+            self.claim(training)
+            return
+        except TrainingSlotTakenError as e:
+            logger.info('%s - waiting', e)
+        while True:
+            await asyncio.sleep(1)
+            try:
+                self.claim(training)
+            except TrainingSlotTakenError:
+                continue
+            logger.info('acquired training slot')
+            return
+
+    @abstractmethod
+    def release(self) -> None:
+        """Give the slot back; a no-op if this slot does not hold it."""
+
+    @property
+    @abstractmethod
+    def held(self) -> bool:
+        """Whether this slot object holds the slot."""
+
+    @property
+    def taken_by_sibling(self) -> bool:
+        return not self.held and self._sibling_holder() is not None
+
+    @abstractmethod
+    def _acquire(self, holder: str) -> None:
+        """Take the slot without waiting; a no-op if this slot already holds it.
+
+        :raise TrainingSlotTakenError: if the slot is held elsewhere, including by a momentary probe.
+        """
+
+    @abstractmethod
+    def _sibling_holder(self) -> str | None:
+        """Who holds the slot, None if it is free or cannot be inspected; only asked while this slot does not hold it."""
+
+    @staticmethod
+    def _describe(training: str) -> str:
+        return f'pid {os.getpid()} on {socket.gethostname()}, {training}'
+
+
+class AlwaysFreeSlot(TrainingSlot):
+    """The slot of a node that has the GPU to itself."""
+
+    def __init__(self) -> None:
+        self._held = False
+
+    def release(self) -> None:
+        self._held = False
+
+    @property
+    def held(self) -> bool:
+        return self._held
+
+    def _acquire(self, holder: str) -> None:
+        self._held = True
+
+    def _sibling_holder(self) -> str | None:
+        return None
+
+
+class FileTrainingSlot(TrainingSlot):
+    """A slot backed by an exclusive `flock` on a file every sibling node mounts.
+
+    The kernel drops the lock with the process, so a node that dies in any way frees the slot.
+    The lock belongs to the inode: the file is only ever truncated and rewritten in place, never
+    replaced, and the mounted path is its directory.
+
+    The file's content names the holder and is diagnostic only — `_sibling_holder` tells a free slot
+    from a taken one by trying the lock, and that probe holds the lock for an instant, so a
+    claim that loses against a probe tries again. The file is made writable for everyone,
+    since sibling containers may run as different users.
+
+    A forked child shares the parent's lock and would keep it alive after the parent releases it,
+    so the child closes its copy right after the fork; the parent's lock is unaffected.
+    """
+
+    _instances: 'weakref.WeakSet[FileTrainingSlot]' = weakref.WeakSet()
+
+    def __init__(self, path: str | Path) -> None:
+        self.path = Path(path)
+        self._fd: int | None = None
+        FileTrainingSlot._instances.add(self)
+
+    def _acquire(self, holder: str) -> None:
+        if self._fd is not None:
+            return
+        fd = self._open()
+        try:
+            self._lock(fd)
+            os.ftruncate(fd, 0)
+            os.write(fd, holder.encode())
+        except BaseException:
+            os.close(fd)
+            raise
+        self._fd = fd
+
+    def release(self) -> None:
+        if self._fd is None:
+            return
+        try:
+            os.ftruncate(self._fd, 0)
+        except OSError:
+            logger.warning('cannot clear the holder from training slot %s', self.path, exc_info=True)
+        finally:
+            os.close(self._fd)
+            self._fd = None
+
+    @property
+    def held(self) -> bool:
+        return self._fd is not None
+
+    def _sibling_holder(self) -> str | None:
+        try:
+            fd = os.open(self.path, os.O_RDONLY)
+        except FileNotFoundError:
+            return None
+        except OSError:
+            logger.warning('cannot inspect training slot %s - treating it as free', self.path, exc_info=True)
+            return None
+        try:
+            self._lock(fd)
+        except TrainingSlotTakenError as e:
+            return e.holder
+        finally:
+            os.close(fd)
+        return None
+
+    @classmethod
+    def _drop_inherited_locks(cls) -> None:
+        for slot in list(cls._instances):
+            if slot._fd is not None:
+                os.close(slot._fd)
+                slot._fd = None
+
+    def _open(self) -> int:
+        fd = os.open(self.path, os.O_RDWR | os.O_CREAT, 0o666)
+        try:
+            os.fchmod(fd, 0o666)
+        except PermissionError:
+            pass
+        return fd
+
+    @staticmethod
+    def _lock(fd: int) -> None:
+        try:
+            fcntl.flock(fd, fcntl.LOCK_EX | fcntl.LOCK_NB)
+        except BlockingIOError as e:
+            raise TrainingSlotTakenError(os.read(fd, 4096).decode(errors='replace') or 'unknown') from e
+
+os.register_at_fork(after_in_child=FileTrainingSlot._drop_inherited_locks)  # pylint: disable=protected-access
+
+
+def training_slot_from_env() -> TrainingSlot:
+    """The slot `TRAINING_SLOT_LOCK` names, or a free one if the variable is unset.
+
+    :raise FileNotFoundError: if the lock file's directory does not exist.
+    """
+    path = os.environ.get(ENV_VAR)
+    if not path:
+        return AlwaysFreeSlot()
+    if not Path(path).parent.is_dir():
+        raise FileNotFoundError(f'{ENV_VAR}={path}: the directory does not exist; mount it from the host')
+    return FileTrainingSlot(path)

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

@@ -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.dist-info → learning_loop_node-0.25.0.dist-info}/RECORD

@@ -22,7 +22,7 @@ learning_loop_node/detector/detector_logic.py,sha256=XJaHNiw5PaHoCWnJ-qsxdUELLev
 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/outbox.py,sha256=HT0GHv6rsQzdbCD8DAT23m1-Q1XEvA_KgwAJLONyX_0,12506
 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
@@ -39,31 +39,32 @@ learning_loop_node/enums/__init__.py,sha256=tjSrhztIQ8W656_QuXfTbbVNtH_wDXP5hpYZ
 learning_loop_node/enums/annotator.py,sha256=mtTAw-8LJIrHcYkBjYHCZuhYEEHS6QzSK8k6BhLusvQ,285
 learning_loop_node/enums/detector.py,sha256=Qvm5LWWR9BfsDxHEQ8YzaPaUuSmp4BescYuV4X4ikwE,512
 learning_loop_node/enums/general.py,sha256=OSpcIpdFjuRByON2xBoiIXl1SgDREIqkwnaxkTFjnHU,164
-learning_loop_node/enums/trainer.py,sha256=VaD63guLO4aKgVfXT0EryPlXKQGegSET3Cp4R0uC_Ew,746
+learning_loop_node/enums/trainer.py,sha256=VdamD-k6Q8kMm_cwPA8iaEUfZvmTxLfsBHQpEi43Ahk,849
 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/entrypoint.py,sha256=EtRMXZKzNapwbIylf4hg8Obw-2EydrLGKUjbIn5H-94,3591
+learning_loop_node/helpers/entrypoint.py,sha256=U4vHFqgke6mXjDL-E4MND8RO7P1k2VBGwdDOOoJvD9Y,3915
 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/misc.py,sha256=zfy1SumqNBwBw6-BYignQJpwMeMf1k_MA3xhYt_GiG4,8939
 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/batch_size.py,sha256=LjA-zyY4hdnSMy6ZsdQOhWbmd5H_DI-9LG3yFInv04E,5476
+learning_loop_node/trainer/cuda.py,sha256=ehpfWqPBCVFu06SnGsjmWbSDRJ8sZc-L_nCtemRLUOU,12567
 learning_loop_node/trainer/downloader.py,sha256=AIyEKM4XYtd6VgmXrP3VayV9DpJzdURK1Brx81ePNSM,1470
-learning_loop_node/trainer/exceptions.py,sha256=IRT7qvn7-RjXaqL8mF21ySm7JWhQtqm1Cl0Gi3DChVY,540
+learning_loop_node/trainer/exceptions.py,sha256=hhd23WvuqV-6fiJo3vCNLN0hmF6TxAH9C5Bc7pCtmUM,940
 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=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/trainer_logic_generic.py,sha256=OtjdZlixz-2UXZkG_o5PaCdXCFvwgGuD5J4I9n5bROY,28812
+learning_loop_node/trainer/trainer_node.py,sha256=GU4LyFA-JC9s4d-PbyKhKWyGMMIzTMqKJFNxM6tDJE4,5316
+learning_loop_node/trainer/training_slot.py,sha256=Njq6Z0T2Lqlr24ftoIp11Pw2OxCo9Z2RGn1tuFhHqLY,6737
 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.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,,
+learning_loop_node-0.25.0.dist-info/METADATA,sha256=UOGaOXyynobyOFSEROGxOPH49odFh2-V2IGZ_LhQNsI,19855
+learning_loop_node-0.25.0.dist-info/WHEEL,sha256=YVMoNqKzERt-wjUZwJ33xBGAwnFl-4cqbYkTtWa4itE,91
+learning_loop_node-0.25.0.dist-info/top_level.txt,sha256=nS5qiFc4stVF_VXVCVJxZnrqVOv3IGX5-_E3sLby6OE,19
+learning_loop_node-0.25.0.dist-info/RECORD,,