diff --git a/AGENTS.md b/AGENTS.md index c7877d5d..b6158d96 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -54,6 +54,15 @@ on 401/429) for the REST API, and a socket.io *client* for status updates and lo - **Annotator** — thin: it forwards the loop frontend's `handle_user_input` events into `AnnotatorLogic` and keeps a per-frontend history. +`trainer/subprocess.py`, `trainer/batch_size.py` and `trainer/metrics.py` hold the parts of a +trainer that are not framework-specific. `iterator_cpu_bound` runs a training generator in a +spawned process and yields its progress through a `maxsize=1` queue, so the event loop stays +responsive, CUDA state stays out of the node process, and the training can never run more than +one item ahead of the bookkeeping. `find_batch_size` probes for the largest power-of-two batch +that fits, around a `fits` predicate the trainer supplies — none of it imports torch, so a node +brings its own way of running a step. `macro_f1` scores the confusion matrix +`_get_new_best_training_state` returns. + `helpers/entrypoint.py` holds what every node's `main.py` repeats: `node_parser` builds a configargparse parser with `--host`/`--port`, `run_node` starts uvicorn. A setting is a flag *and* an environment variable from one declaration — `--conf-threshold` reads diff --git a/learning_loop_node/tests/unit/test_batch_size.py b/learning_loop_node/tests/unit/test_batch_size.py new file mode 100644 index 00000000..6b499a6a --- /dev/null +++ b/learning_loop_node/tests/unit/test_batch_size.py @@ -0,0 +1,118 @@ +from collections.abc import Callable + +import pytest + +from ...trainer.batch_size import ( + MIN_TRAIN_STEPS_PER_EPOCH, + batch_count, + dataset_limit, + find_batch_size, + is_out_of_memory, + no_gpu_batch_size, + smaller_pot, +) + + +def _recording(capacity: int) -> tuple[Callable[[int], bool], list[int]]: + """A `fits` predicate for a machine of `capacity`, plus the sizes it gets asked about.""" + calls: list[int] = [] + + def fits(batch_size: int) -> bool: + calls.append(batch_size) + return batch_size <= capacity + + return fits, calls + + +def _fits_up_to(capacity: int) -> Callable[[int], bool]: + fits, _ = _recording(capacity) + return fits + + +def test_the_search_doubles_up_to_the_limit(): + fits, calls = _recording(1024) + assert find_batch_size(fits, limit=16) == 16 + assert calls == [1, 2, 4, 8, 16] + + +def test_the_search_backs_off_to_the_last_size_that_fit(): + fits, calls = _recording(20) + assert find_batch_size(fits, limit=64) == 16 + assert calls == [1, 2, 4, 8, 16, 32] # probes 32, fails, keeps 16 + + +def test_a_limit_that_is_not_a_power_of_two_is_rounded_down(): + assert find_batch_size(_fits_up_to(1024), limit=48) == 32 + assert find_batch_size(_fits_up_to(1024), limit=1) == 1 + + +def test_a_machine_that_cannot_take_one_sample_is_an_error(): + fits, calls = _recording(0) + with pytest.raises(RuntimeError, match='batch size 1 does not fit'): + find_batch_size(fits, limit=64) + assert calls == [1], 'must give up instead of probing larger sizes' + + +@pytest.mark.parametrize('capacity', range(1, 130)) +def test_the_result_is_always_the_largest_power_of_two_that_fits(capacity: int): + assert find_batch_size(_fits_up_to(capacity), limit=512) == smaller_pot(capacity) + + +def test_equal_hardware_yields_an_equal_recipe(): + """Only powers of two, so two machines of similar size train identically.""" + assert find_batch_size(_fits_up_to(37), limit=512) == find_batch_size(_fits_up_to(39), limit=512) + + +def test_smaller_pot(): + assert [smaller_pot(n) for n in (1, 2, 3, 4, 7, 8, 15, 1293)] == [1, 2, 2, 4, 4, 8, 8, 1024] + with pytest.raises(ValueError, match='n must be >= 1'): + smaller_pot(0) + + +def test_batch_count_covers_the_whole_set_without_overshooting_by_a_batch(): + for sample_count in (1, 7, 8, 900, 5000): + for batch_size in (1, 2, 8, 64, 512): + covered = batch_count(sample_count, batch_size) * batch_size + assert covered >= sample_count + assert covered - sample_count < batch_size + + +def test_the_dataset_limit_keeps_enough_steps_per_epoch(): + for sample_count in (8, 20, 47, 100, 1000, 118_000): + limit = smaller_pot(dataset_limit(sample_count)) + assert sample_count // limit >= MIN_TRAIN_STEPS_PER_EPOCH + + +def test_the_dataset_limit_stays_usable_for_a_tiny_set(): + assert dataset_limit(7) == 1 + with pytest.raises(ValueError, match='sample_count must be >= 1'): + dataset_limit(0) + + +def test_memory_still_decides_below_the_dataset_limit(): + """The bound is a ceiling only: a card that fits just 2 keeps training at 2.""" + assert find_batch_size(_fits_up_to(2), limit=dataset_limit(20)) == 2 + + +def test_without_a_gpu_the_fallback_respects_the_limit(): + assert no_gpu_batch_size(1024, 'probe') == 8 + assert no_gpu_batch_size(2, 'probe') == 2 + + +@pytest.mark.parametrize('message', [ + 'CUDA out of memory. Tried to allocate 20.00 MiB', + 'cuDNN error: CUDNN_STATUS_ALLOC_FAILED', + 'CUDA error: unknown error', +]) +def test_allocation_failures_are_recognised_however_they_surface(message: str): + assert is_out_of_memory(RuntimeError(message)) + + +def test_a_real_bug_is_not_mistaken_for_a_full_card(): + """A trainer catching bare RuntimeError treats every crash as 'too big'; this does not.""" + assert not is_out_of_memory(RuntimeError('shape mismatch in forward pass')) + assert not is_out_of_memory(ValueError('bad config')) + + +def test_the_host_running_out_of_memory_counts_too(): + assert is_out_of_memory(MemoryError()) diff --git a/learning_loop_node/tests/unit/test_metrics.py b/learning_loop_node/tests/unit/test_metrics.py new file mode 100644 index 00000000..7c8537dc --- /dev/null +++ b/learning_loop_node/tests/unit/test_metrics.py @@ -0,0 +1,39 @@ +import pytest + +from ...trainer.metrics import category_f1, confusion_matrix_from_counts, macro_f1 + + +def test_the_score_averages_the_categories_instead_of_pooling_them(): + """Counts from a real epoch: face 248/19/59 and hand 39/54/73 - 74% pooled, 62% averaged.""" + assert macro_f1({'face': {'tp': 248, 'fp': 19, 'fn': 59}, + 'hand': {'tp': 39, 'fp': 54, 'fn': 73}}) == pytest.approx(0.622, abs=0.001) + assert macro_f1({'both': {'tp': 287, 'fp': 73, 'fn': 132}}) == pytest.approx(0.737, abs=0.001) + + +def test_a_rare_category_carries_the_same_weight(): + """A category the model never finds halves the score, however few instances it has.""" + assert macro_f1({'frequent': {'tp': 1000, 'fp': 0, 'fn': 0}, + 'rare': {'tp': 0, 'fp': 0, 'fn': 3}}) == pytest.approx(0.5) + + +def test_a_category_without_any_counts_scores_zero_instead_of_raising(): + assert macro_f1({'unseen': {'tp': 0, 'fp': 0, 'fn': 0}}) == 0.0 + assert category_f1({'tp': 0, 'fp': 0, 'fn': 0}) == 0.0 + + +def test_an_empty_confusion_matrix_scores_zero(): + assert macro_f1({}) == 0.0 + + +def test_a_perfect_category_scores_one(): + assert category_f1({'tp': 10, 'fp': 0, 'fn': 0}) == pytest.approx(1.0) + + +def test_precision_and_recall_are_balanced(): + assert category_f1({'tp': 5, 'fp': 5, 'fn': 0}) == category_f1({'tp': 5, 'fp': 0, 'fn': 5}) + + +def test_counts_can_be_assembled_from_separate_mappings(): + matrix = confusion_matrix_from_counts({'a': 3}, {'a': 1, 'b': 2}, {'b': 4}) + assert matrix == {'a': {'tp': 3, 'fp': 1, 'fn': 0}, + 'b': {'tp': 0, 'fp': 2, 'fn': 4}} diff --git a/learning_loop_node/tests/unit/test_subprocess.py b/learning_loop_node/tests/unit/test_subprocess.py new file mode 100644 index 00000000..f2517ced --- /dev/null +++ b/learning_loop_node/tests/unit/test_subprocess.py @@ -0,0 +1,54 @@ +import pytest + +from ...trainer.subprocess import iterator_cpu_bound + + +def counting(limit: int): + yield from range(limit) + + +def raising_after(limit: int): + yield from range(limit) + raise RuntimeError('the training crashed') + + +def empty(): + return iter(()) + + +async def _collect(iterator) -> list: + return [item async for item in iterator] + + +async def test_everything_the_generator_yields_arrives_in_order(): + async with iterator_cpu_bound(counting, 5) as iterator: + assert await _collect(iterator) == [0, 1, 2, 3, 4] + + +async def test_a_generator_that_yields_nothing_simply_finishes(): + async with iterator_cpu_bound(empty) as iterator: + assert await _collect(iterator) == [] + + +async def test_a_failure_in_the_process_is_raised_in_the_caller(): + """Otherwise a crashed training would look like a training that finished.""" + with pytest.raises(RuntimeError, match='the training crashed'): + async with iterator_cpu_bound(raising_after, 2) as iterator: + await _collect(iterator) + + +async def test_what_the_generator_produced_before_failing_still_arrives(): + received = [] + with pytest.raises(RuntimeError): + async with iterator_cpu_bound(raising_after, 3) as iterator: + async for item in iterator: + received.append(item) + assert received == [0, 1, 2] + + +async def test_leaving_early_does_not_leave_the_process_running(): + async with iterator_cpu_bound(counting, 1000) as iterator: + async for item in iterator: + if item == 2: + break + # the context manager killed and joined the process; reaching here without hanging is the test diff --git a/learning_loop_node/trainer/batch_size.py b/learning_loop_node/trainer/batch_size.py new file mode 100644 index 00000000..707bb8c6 --- /dev/null +++ b/learning_loop_node/trainer/batch_size.py @@ -0,0 +1,86 @@ +"""Choosing a batch size by probing, rather than configuring one. + +The largest batch that fits depends on the model, the image resolution and the card, so a +configured value is either wrong on some machines or too small on all of them. Probing answers +it directly: run a representative step at doubling sizes and keep the last one that survived. + +What "a representative step" means is the trainer's business — it supplies a ``fits`` predicate +that runs a real step and reports whether it ran out of memory. Everything here is that +predicate's scaffolding, and none of it imports a deep-learning framework. + +Only powers of two are visited, so equal hardware yields an equal recipe. That matters when a +trainer scales its learning rate by the batch size: a probe that returned 37 on one machine and +39 on another would make the two trainings quietly different. + +Adapted from PyTorch Lightning's ``BatchSizeFinder`` (power-scaling mode). +Copyright The Lightning AI team. Licensed under the Apache License, Version 2.0. +https://github.com/Lightning-AI/pytorch-lightning +""" + +import logging +from collections.abc import Callable + +logger = logging.getLogger(__name__) + +NO_GPU_BATCH_SIZE = 8 +"""What to fall back on with no GPU to probe: enough to make progress, small enough to fit.""" + +MIN_TRAIN_STEPS_PER_EPOCH = 8 +"""Below this an epoch is one or two optimizer steps, which trains poorly whatever the card.""" + + +def find_batch_size(fits: Callable[[int], bool], *, limit: int) -> int: + """Return the largest power-of-two batch size that fits, never exceeding ``limit``. + + :param fits: Runs a representative probe; ``False`` on out-of-memory. + :raises RuntimeError: If not even a batch size of 1 fits. + """ + limit = smaller_pot(limit) + if not fits(1): + raise RuntimeError('batch size 1 does not fit in memory') + + size = 1 + while size < limit and fits(size * 2): + size *= 2 + + return size + + +def dataset_limit(sample_count: int) -> int: + """The batch size ceiling that still leaves ``MIN_TRAIN_STEPS_PER_EPOCH`` steps per epoch.""" + if sample_count < 1: + raise ValueError(f'sample_count must be >= 1, got {sample_count}') + return max(1, sample_count // MIN_TRAIN_STEPS_PER_EPOCH) + + +def no_gpu_batch_size(limit: int, probe: str) -> int: + """The batch size to fall back on when there is no GPU to probe.""" + batch_size = min(smaller_pot(limit), NO_GPU_BATCH_SIZE) + logger.warning('%s: CUDA is unavailable; using batch size %d without probing', probe, batch_size) + return batch_size + + +def smaller_pot(n: int) -> int: + """The largest power of two that is <= ``n``.""" + if n < 1: + raise ValueError(f'n must be >= 1, got {n}') + return 1 << (n.bit_length() - 1) + + +def batch_count(sample_count: int, batch_size: int) -> int: + """How many batches a loader yields, the last one possibly short.""" + return -(-sample_count // batch_size) + + +def is_out_of_memory(exception: BaseException) -> bool: + """Whether the exception signals exhausted memory, on the GPU or the host. + + Allocation failures do not all surface as a framework's dedicated error type: cuDNN and + cuBLAS workspaces raise a plain ``RuntimeError``. Catching those by message is what keeps a + probe from mistaking a real bug for a full card — a trainer that catches bare + ``RuntimeError`` around its probe silently treats every crash as "too big". + """ + if isinstance(exception, MemoryError): + return True + message = str(exception).lower() + return any(text in message for text in ('out of memory', 'alloc_failed', 'cuda error: unknown error')) diff --git a/learning_loop_node/trainer/metrics.py b/learning_loop_node/trainer/metrics.py new file mode 100644 index 00000000..aee553b8 --- /dev/null +++ b/learning_loop_node/trainer/metrics.py @@ -0,0 +1,36 @@ +"""Scoring a training from the confusion matrix the loop stores. + +The loop keeps one ``{'tp': .., 'fp': .., 'fn': ..}`` per category, which is what +``TrainerLogicGeneric._get_new_best_training_state`` returns. Deciding whether an epoch beat +the previous best means reducing that to one number, and every trainer needs the same one. +""" + +import statistics +from collections.abc import Mapping + + +def macro_f1(confusion_matrix: Mapping[str, Mapping[str, int]]) -> float: + """The unweighted mean of the per-category F1 scores, as the loop's UI shows it by default. + + Unweighted is the point: a rare category the model never finds drags the score down as much + as a frequent one, where pooling the counts first would hide it. + """ + scores = [category_f1(counts) for counts in confusion_matrix.values()] + return statistics.mean(scores) if scores else 0.0 + + +def category_f1(counts: Mapping[str, int]) -> float: + """The F1 score of one category, 0.0 where it is undefined rather than a division error.""" + tp, fp, fn = counts['tp'], counts['fp'], counts['fn'] + precision = tp / (tp + fp) if tp + fp else 0.0 + recall = tp / (tp + fn) if tp + fn else 0.0 + return 2 * precision * recall / (precision + recall) if precision + recall else 0.0 + + +def confusion_matrix_from_counts(true_positives: Mapping[str, int], false_positives: Mapping[str, int], + false_negatives: Mapping[str, int]) -> dict[str, dict[str, int]]: + """Assemble the loop's confusion matrix from three per-category count mappings.""" + return {category: {'tp': true_positives.get(category, 0), + 'fp': false_positives.get(category, 0), + 'fn': false_negatives.get(category, 0)} + for category in {*true_positives, *false_positives, *false_negatives}} diff --git a/learning_loop_node/trainer/subprocess.py b/learning_loop_node/trainer/subprocess.py new file mode 100644 index 00000000..dbfa181b --- /dev/null +++ b/learning_loop_node/trainer/subprocess.py @@ -0,0 +1,93 @@ +"""Run a blocking, CPU-bound generator in its own process without blocking the event loop. + +A trainer that trains in-process has a problem: the training must not stall the node, and CUDA +state must stay out of the node process so a crashed training cannot take the node with it. +:func:`iterator_cpu_bound` runs the generator in a spawned process and yields what it produces +through a ``maxsize=1`` queue, so the producer can never run more than one item ahead of the +bookkeeping that consumes it — which is what lets a trainer alternate between two model files +and know the one it is copying is not being rewritten. + +Exceptions raised inside the process are re-raised in the caller, and the process is killed if +the caller leaves the context early. +""" +from __future__ import annotations + +import asyncio +import multiprocessing +import queue +import traceback +from collections.abc import AsyncGenerator, Callable, Iterator +from contextlib import asynccontextmanager +from multiprocessing.queues import Queue as MPQueue +from typing import Any, ParamSpec, TypeVar + +T = TypeVar('T') +P = ParamSpec('P') + + +class IteratorDone: + pass + + +def _iterator_wrapper( + it: Callable[..., Iterator[T]], + state_queue: MPQueue[T | Exception | IteratorDone], + args: tuple[Any, ...], + kwargs: dict[str, Any], +) -> None: + try: + for data in it(*args, **kwargs): + state_queue.put(data) + except Exception as e: + print(traceback.format_exc()) + state_queue.put(e) + + state_queue.put(IteratorDone()) + + +async def _iterator_cpu_bound_inner( + it: Callable[P, Iterator[T]], + *args: P.args, + **kwargs: P.kwargs, +) -> AsyncGenerator[T, None]: + state_queue: MPQueue[T | Exception | IteratorDone] = multiprocessing.Queue(maxsize=1) + process = multiprocessing.Process( + target=_iterator_wrapper, + args=(it, state_queue, args, kwargs), + name='iterator_cpu_bound', + ) + + process.start() + + try: + while True: + try: + item = await asyncio.to_thread(state_queue.get, True, 0.5) + except queue.Empty: + if not process.is_alive(): + break + continue + match item: + case IteratorDone(): + break + case Exception() as e: + raise e + case _ as other: + yield other + finally: + if process.is_alive(): + process.kill() + process.join() + + +@asynccontextmanager +async def iterator_cpu_bound( + it: Callable[P, Iterator[T]], + *args: P.args, + **kwargs: P.kwargs, +) -> AsyncGenerator[AsyncGenerator[T, None], None]: + iterator = _iterator_cpu_bound_inner(it, *args, **kwargs) + try: + yield iterator + finally: + await asyncio.shield(iterator.aclose()) diff --git a/learning_loop_node/trainer/trainer_logic.py b/learning_loop_node/trainer/trainer_logic.py index c5a9bade..9827b41b 100644 --- a/learning_loop_node/trainer/trainer_logic.py +++ b/learning_loop_node/trainer/trainer_logic.py @@ -178,9 +178,19 @@ async def _resume(self) -> None: '''Is called when self.can_resume() returns True. One may resume the training on a previously trained model stored by self.on_model_published(basic_model).''' - @abstractmethod def _get_executor_error_from_log(self) -> Optional[str]: - '''Should be used to provide error informations to the Learning Loop by extracting data from self.executor.get_log().''' + '''Reports what went wrong to the Learning Loop by reading self.executor's log. + + The default recognises the CUDA failures every trainer hits. Override to add messages a + particular training framework produces, and call super() to keep these.''' + if self._executor is None: + return None + for line in self._executor.get_log_by_lines(tail=50): + if 'CUDA out of memory' in line: + return 'graphics card is out of memory' + if 'CUDA error: invalid device ordinal' in line: + return 'graphics card not found' + return None @abstractmethod async def _detect(self, model_information: ModelInformation, images: List[str], model_folder: str) -> List[Detections]: