diff --git a/backend/app/services/ml/ccip.py b/backend/app/services/ml/ccip.py index 833b55c..fc80ec2 100644 --- a/backend/app/services/ml/ccip.py +++ b/backend/app/services/ml/ccip.py @@ -35,6 +35,58 @@ DEFAULT_SIM_THRESHOLD = 0.85 _FIGURE_KINDS = ("face", "figure") +# How many cosine scores to hold in memory at once, per matmul block. +# 4M float32 is 16 MB — small enough to stay in cache-friendly territory on the +# shared ml lane, large enough that the per-call overhead stops mattering. +_MAX_SCORE_ELEMS = 4_000_000 + + +def char_maxima(q_by_image, allref, seg, np, *, max_elems=_MAX_SCORE_ELEMS): + """(n_images, n_chars) — each image's best cosine to each character. + + `q_by_image` is one L2-normalised `(n_figures, dim)` array per image, in + the order the answer comes back in. `allref` is every character's + prototypes stacked, and `seg` their per-character start offsets into it. + + ## Why this is batched, and why that is safe + + `scheduled_ccip_auto_apply` did this one image at a time — a `(nq, dim) @ + (dim, total)` product per image, over every image in the library on every + run. At ~119k images that is 119k separate matmuls, each too small to pay + for its own BLAS setup, and on 2026-09-23 the daily sweep hit its 1800s + soft limit on the operator's instance. + + Batching changes no arithmetic. The score a character gets for an image is + a max over that image's figures AND over that character's prototypes, and + max does not care in what order or grouping it is taken — so reducing the + prototype axis first (per row, inside a block) and the figure axis after + (per image, across blocks) gives exactly what the per-image loop gave. + That equivalence is what `test_char_maxima_matches_the_per_image_loop` + pins, against the naive form written out longhand. + + Blocked by ROWS rather than done in one product, because the full score + matrix is (all figures in the chunk x every prototype) and that grows with + the library on both axes. The block bound is on elements, so the memory + this uses stays flat as either axis grows. + """ + counts = [len(q) for q in q_by_image] + rows = np.vstack(q_by_image) + total = max(int(allref.shape[0]), 1) + block = max(1, max_elems // total) + + per_row = np.empty((rows.shape[0], len(seg)), dtype=np.float32) + for a in range(0, rows.shape[0], block): + scores = rows[a:a + block] @ allref.T + per_row[a:a + block] = np.maximum.reduceat(scores, seg, axis=1) + + # Start offset of each image's rows. Every image has at least one figure — + # it is in `q_by_image` because a region produced it — so these strictly + # increase, which is what `reduceat` needs to reduce rather than pass a row + # through untouched. + starts = np.cumsum([0] + counts[:-1]) + return np.maximum.reduceat(per_row, starts, axis=0) + + async def _settings_threshold(session: AsyncSession) -> float: val = ( await session.execute( diff --git a/backend/app/services/ml/embedder.py b/backend/app/services/ml/embedder.py index d55646c..abb87cf 100644 --- a/backend/app/services/ml/embedder.py +++ b/backend/app/services/ml/embedder.py @@ -11,12 +11,21 @@ from pathlib import Path import numpy as np from PIL import Image, ImageFile +from ..worker_lanes import LANES_BY_NAME + ImageFile.LOAD_TRUNCATED_IMAGES = True # Cap torch's intra-op threads so each ml-worker replica is a bounded core -# consumer on a shared node (torch otherwise uses all cores). Keep -# N_replicas × this within the cores allotted to ML to avoid oversubscription. -_INTRA_OP_THREADS = 4 +# consumer on a shared node (torch otherwise uses all cores). +# +# Read from the lane rather than restated here. This was a literal 4 beside a +# comment reading "keep N_replicas x this within the cores allotted to ML" — +# a constraint written where nothing could act on it, and nothing did: the ML +# ceiling came from memory alone, offered the operator ~49 slots on a +# large-memory host, and the lane spent 2026-09-23 with ~200 torch threads on +# it. `derived_ceiling` now divides the cores by this number, which only means +# anything while the two are the same number. +_INTRA_OP_THREADS = LANES_BY_NAME["ml"].threads_per_slot DEFAULT_MODEL_NAME = os.environ.get( "SIGLIP_MODEL_NAME", "google/siglip-so400m-patch14-384" diff --git a/backend/app/services/worker_lanes.py b/backend/app/services/worker_lanes.py index 8a23513..aad8f9a 100644 --- a/backend/app/services/worker_lanes.py +++ b/backend/app/services/worker_lanes.py @@ -139,8 +139,22 @@ class Lane: # tells them to raise it rather than quietly consuming the machine. default_slots_cap: int # True when a slot costs a copy of the ML model rather than just a process. - # The only lane whose ceiling is decided by memory instead of by cores. + # Such a lane is bounded by memory AS WELL AS by cores, never instead of. memory_bound: bool = False + # CPU threads ONE slot uses. More than one for a lane whose work is an + # inference library with its own thread pool: `services/ml/embedder.py` + # calls `torch.set_num_threads` with this number, so a slot is four cores' + # worth of demand rather than one process's. + # + # It lives here because the CEILING has to know it. It was a private + # constant in the embedder with a comment saying "keep N_replicas x this + # within the cores allotted to ML" — a rule stated where nothing could + # enforce it. Nothing did: the ML ceiling was computed from memory alone, + # so a large-memory host offered ~49 slots, the operator took them, and + # 2026-09-23's log shows ~200 torch threads fighting over the box — + # embeds at 107-246s each, and the daily CCIP sweep sharing that pool + # timing out at 1800s. + threads_per_slot: int = 1 # Models this lane downloads the first time it is enabled. Empty for every # lane that needs none, which is how the UI knows whether to warn at all. models: tuple[ModelRequirement, ...] = () @@ -209,6 +223,7 @@ LANES: tuple[Lane, ...] = ( entrypoint_role="ml-worker", default_slots_cap=0, memory_bound=True, + threads_per_slot=4, models=(SIGLIP_MODEL,), optional=True, ), @@ -370,6 +385,21 @@ def lane_for_node(hostname: str) -> Lane | None: return LANES_BY_NAME.get(hostname.split("@", 1)[0]) +def _cpu_bound_slots(lane: Lane) -> int: + """How many slots this container's cores can feed, at `threads_per_slot`. + + Never zero: a machine with fewer cores than one slot wants still runs the + lane, just slowly. That is a real trade an operator may want, and refusing + to offer the lane at all on a small box would make ML unreachable there — + unlike the memory bound, where the honest answer IS zero, because the + first task would OOM the container rather than merely be slow. + """ + cores = container_cpu_count() + if cores is None: + return UNKNOWN_CEILING + return max(MIN_CEILING, cores // lane.threads_per_slot) + + def derived_ceiling(lane: Lane) -> int: """The most slots `lane` may be given on this container. @@ -377,26 +407,31 @@ def derived_ceiling(lane: Lane) -> int: is bounded by what it has NOW rather than by what it had when its row was written. """ - if lane.memory_bound: - total = container_memory_bytes() - if total is None: - log.warning( - "worker_lanes: cannot read a memory limit; capping %s at %d", - lane.name, UNKNOWN_CEILING, - ) - return UNKNOWN_CEILING - usable = total - RESERVED_BYTES - if usable < ML_BYTES_PER_SLOT: - # Honestly zero. A box that cannot hold one model alongside the web - # process must be told it cannot run tagging, not sold a slot that - # will OOM the container the first time it is used. - return 0 - return int(usable // ML_BYTES_PER_SLOT) + by_cpu = _cpu_bound_slots(lane) + if not lane.memory_bound: + return by_cpu - cores = container_cpu_count() - if cores is None: + total = container_memory_bytes() + if total is None: + log.warning( + "worker_lanes: cannot read a memory limit; capping %s at %d", + lane.name, UNKNOWN_CEILING, + ) return UNKNOWN_CEILING - return max(MIN_CEILING, cores) + usable = total - RESERVED_BYTES + if usable < ML_BYTES_PER_SLOT: + # Honestly zero. A box that cannot hold one model alongside the web + # process must be told it cannot run tagging, not sold a slot that + # will OOM the container the first time it is used. + return 0 + + # BOTH bounds, whichever binds first. Memory alone was the whole answer + # until 2026-09-23, and on a large-memory host that is the wrong one: RAM + # said ~49 slots, and each of those slots wants `threads_per_slot` cores. + # The operator raised the cap to what the dial offered and the lane + # starved itself — a control is not allowed to offer a number the machine + # cannot feed. + return min(int(usable // ML_BYTES_PER_SLOT), by_cpu) def ceilings() -> dict[str, int]: diff --git a/backend/app/tasks/ml.py b/backend/app/tasks/ml.py index 4388820..e21c1aa 100644 --- a/backend/app/tasks/ml.py +++ b/backend/app/tasks/ml.py @@ -488,7 +488,7 @@ def scheduled_ccip_auto_apply() -> str: from ..models import ImageRegion, MLSettings, Tag, TagKind from ..models.tag import image_tag - from ..services.ml.ccip import _FIGURE_KINDS + from ..services.ml.ccip import _FIGURE_KINDS, char_maxima from ..services.ml.training_data import _applied_or_rejected, _l2norm SessionLocal = _sync_session_factory() @@ -553,11 +553,22 @@ def scheduled_ccip_auto_apply() -> str: by_img: dict[int, list] = {} for iid, vec in rows: by_img.setdefault(iid, []).append(vec) - for iid, vecs in by_img.items(): - q = _l2norm(np.asarray(vecs, dtype=np.float32), np) # (nq, 768) - colmax = (q @ allref.T).max(axis=0) # (total,) - charmax = np.maximum.reduceat(colmax, seg) # (n_chars,) - for ci in np.where(charmax >= thr)[0]: + if not by_img: + continue + + # One matmul per BLOCK of figures, not one per image. This loop ran + # over every image in the library on every daily run and did a + # matmul too small to pay for itself each time; it hit the 1800s + # soft limit on the operator's instance on 2026-09-23. Same + # arithmetic — see `char_maxima`. + iids = list(by_img) + charmax = char_maxima( + [_l2norm(np.asarray(by_img[i], dtype=np.float32), np) for i in iids], + allref, seg, np, + ) # (n_img, n_chars) + + for row, iid in enumerate(iids): + for ci in np.where(charmax[row] >= thr)[0]: t = ref_tags[int(ci)] if iid in skip[t]: continue diff --git a/frontend/src/components/settings/SystemHealthTab.vue b/frontend/src/components/settings/SystemHealthTab.vue index 4ae4b34..4b96486 100644 --- a/frontend/src/components/settings/SystemHealthTab.vue +++ b/frontend/src/components/settings/SystemHealthTab.vue @@ -170,9 +170,9 @@ FabledCurator to work at all; ML tagging is the one that is genuinely optional, and it ships at zero because switching it on downloads a model. Up to N beneath each dial is what this machine can - hold — memory for ML tagging, processor cores for the rest — and it is - recalculated from the container's real limits every time this page - loads. + hold — processor cores for every lane, and for ML tagging whichever + runs out first, its cores or its memory. It is recalculated from the + container's real limits every time this page loads.

The shipped caps are one of each, which is right for a first boot and @@ -218,8 +218,9 @@

- Each worker loads its own copy, which is why this machine allows it - at most {{ lane.ceiling }}. + Each worker loads its own copy and asks for several processor cores + while it runs, which is why this machine allows it at most + {{ lane.ceiling }}. diff --git a/tests/test_ml_dry_helpers.py b/tests/test_ml_dry_helpers.py index bff8d27..67bed9d 100644 --- a/tests/test_ml_dry_helpers.py +++ b/tests/test_ml_dry_helpers.py @@ -57,3 +57,88 @@ def test_applied_or_rejected_unions_applied_any_source_and_rejected(db_sync): assert skip[b.id] == {imgs[3].id} assert imgs[4].id not in skip[a.id] assert imgs[4].id not in skip[b.id] + + +# --- the CCIP auto-apply sweep's scorer --------------------------------------- +# +# `scheduled_ccip_auto_apply` scored one image per matmul, over every image in +# the library, on every daily run — and on 2026-09-23 it hit its 1800s soft +# limit on the operator's instance. `char_maxima` does the same arithmetic in +# blocks. These pin THAT: same answer, whatever the blocking. + + +def _score_fixture(np): + """Four images with 1-3 figures each, three characters with 2/5/1 + prototypes. Deliberately ragged — equal group sizes would let a wrong + `reduceat` offset pass.""" + from backend.app.services.ml.training_data import _l2norm + + rng = np.random.default_rng(7) + dim = 16 + q_by_image = [ + _l2norm(rng.standard_normal((n, dim)).astype(np.float32), np) + for n in (1, 3, 2, 1) + ] + mats = [ + _l2norm(rng.standard_normal((k, dim)).astype(np.float32), np) + for k in (2, 5, 1) + ] + allref = np.vstack(mats) + seg = np.cumsum([0] + [len(m) for m in mats])[:-1] + return q_by_image, allref, seg + + +def _naive(q_by_image, allref, seg, np): + """The loop as it was written before batching, kept longhand. The point of + comparing against this rather than against a stored array is that it is + the OLD CODE — if the batched form ever diverges, this says so in the + terms the change was justified in.""" + return np.vstack([ + np.maximum.reduceat((q @ allref.T).max(axis=0), seg) for q in q_by_image + ]) + + +def test_char_maxima_matches_the_per_image_loop(): + import numpy as np + + from backend.app.services.ml.ccip import char_maxima + + q_by_image, allref, seg = _score_fixture(np) + got = char_maxima(q_by_image, allref, seg, np) + + assert got.shape == (len(q_by_image), len(seg)) + np.testing.assert_allclose( + got, _naive(q_by_image, allref, seg, np), rtol=1e-6, atol=1e-6, + ) + + +def test_the_answer_does_not_depend_on_where_the_blocks_fall(): + """The one thing batching could get wrong. Rows are reduced over the + PROTOTYPE axis inside a block and over the FIGURE axis afterwards, so a + block boundary may fall in the middle of an image's figures — which is + safe only because max does not care how it is grouped. `max_elems=1` + forces a boundary between every single row.""" + import numpy as np + + from backend.app.services.ml.ccip import char_maxima + + q_by_image, allref, seg = _score_fixture(np) + whole = char_maxima(q_by_image, allref, seg, np, max_elems=10_000_000) + split = char_maxima(q_by_image, allref, seg, np, max_elems=1) + + np.testing.assert_allclose(whole, split, rtol=1e-6, atol=1e-6) + + +def test_one_character_and_one_figure_still_reduces(): + """The degenerate shape `reduceat` is easiest to get wrong: a single + segment starting at 0, and a single row.""" + import numpy as np + + from backend.app.services.ml.ccip import char_maxima + + q = np.array([[1.0, 0.0]], dtype=np.float32) + allref = np.array([[1.0, 0.0], [0.0, 1.0]], dtype=np.float32) + got = char_maxima([q], allref, np.array([0]), np) + + assert got.shape == (1, 1) + assert got[0][0] == pytest.approx(1.0) diff --git a/tests/test_worker_lanes.py b/tests/test_worker_lanes.py index 160a552..cdf0a2f 100644 --- a/tests/test_worker_lanes.py +++ b/tests/test_worker_lanes.py @@ -153,10 +153,60 @@ def test_v1_sentinel_is_recognised_as_unlimited(monkeypatch, tmp_path): def test_ml_ceiling_is_memory_divided_by_per_slot_cost(monkeypatch, tmp_path): # 2 GiB reserved for web and the other lanes, then 4 GiB per model copy. + # Cores pinned high so the memory bound is the one being read here. _point_memory_at(monkeypatch, tmp_path, str(14 * wl.GIB)) + monkeypatch.setattr(wl, "container_cpu_count", lambda: 32) assert wl.derived_ceiling(wl.LANES_BY_NAME["ml"]) == 3 +def test_ml_is_bounded_by_cores_as_well_as_memory(monkeypatch, tmp_path): + """The bug of 2026-09-23, in one assertion. + + The ML ceiling was memory ALONE. On the operator's large-memory host that + offered ~49 slots; they took them, and each slot asks torch for + `threads_per_slot` cores — so the lane ran ~200 threads over a box that + has nowhere near that many. Embeds that should be seconds took 107-246s, + and the daily CCIP sweep sharing that pool died on its 1800s soft limit. + + Memory here says 49. The cores say 8 / 4 = 2, and the smaller bound is the + only honest one: a control must not offer a number the machine cannot + feed. + """ + _point_memory_at(monkeypatch, tmp_path, str(200 * wl.GIB)) + monkeypatch.setattr(wl, "container_cpu_count", lambda: 8) + + ml = wl.LANES_BY_NAME["ml"] + assert ml.threads_per_slot == 4 + assert wl.derived_ceiling(ml) == 2 + + +def test_a_few_cores_still_offer_one_ml_slot_rather_than_none( + monkeypatch, tmp_path, +): + """The two bounds fail differently, deliberately. Too little MEMORY is + honestly zero — the first task would OOM the container. Too few CORES is + merely slow, which is a trade an operator may want, so it floors at one + rather than making the lane unreachable on a small box.""" + _point_memory_at(monkeypatch, tmp_path, str(200 * wl.GIB)) + monkeypatch.setattr(wl, "container_cpu_count", lambda: 1) + assert wl.derived_ceiling(wl.LANES_BY_NAME["ml"]) == 1 + + +def test_the_embedder_asks_for_exactly_what_the_ceiling_budgeted(monkeypatch): + """The two halves of the same number, tied. + + `threads_per_slot` is only meaningful because `embedder.load()` calls + `torch.set_num_threads` with it. It lived in the embedder as a private 4 + beside a comment saying "keep N_replicas x this within the cores allotted + to ML" — a constraint stated where nothing could enforce it, and nothing + did. If these two ever drift, the ceiling is budgeting cores for a demand + the worker does not make, and nothing else would notice. + """ + from backend.app.services.ml import embedder + + assert embedder._INTRA_OP_THREADS == wl.LANES_BY_NAME["ml"].threads_per_slot + + def test_a_small_box_is_told_it_cannot_run_tagging(monkeypatch, tmp_path): """Honestly zero rather than a floor of one. A 4GB box cannot hold a model alongside the web process, and offering a slot that OOMs the container the @@ -206,6 +256,7 @@ def test_the_ceiling_is_computed_not_stored(monkeypatch, tmp_path): ceiling when the container's limits change, with no row edit. A stored ceiling would keep authorising what the box no longer has.""" lane = wl.LANES_BY_NAME["ml"] + monkeypatch.setattr(wl, "container_cpu_count", lambda: 32) _point_memory_at(monkeypatch, tmp_path, str(34 * wl.GIB)) big = wl.derived_ceiling(lane) (tmp_path / "memory.max").write_text(str(10 * wl.GIB))