Share the framework-agnostic parts of a trainer

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

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())

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}}

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

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'))

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}}