3614b23a2c
The server segfaulted a few seconds into every live scan that used the crop lane. The crash report put the fault in libtorch's Metal shader library (`MetalShaderLibrary::exec_unary_kernel` inside an unordered_map insert), reached from `copy_cast_kernel_mps` under `clone`/`reshape`. Reproduced deterministically: eight threads calling `embed()` on a fresh process crash it, while the same eight threads after one warm forward are fine. Metal's shader library is process-global and is not thread-safe while kernels are still being compiled into it, which is exactly what the first forward does. The detector already serializes its predict for this reason; the gallery did not. Two changes: * A module-level `_MPS_INFER_LOCK` covers the whole MPS region of a forward - model load, `.to(device)`, the forward, and the normalize that follows it. Normalize dispatches its own Metal kernels, so leaving it outside the lock left the same first-compile race open. Module-level because the state being protected is the process's, not one gallery's. * A cache hit skips the embedding forward, so the model used to be built and the kernels compiled by whichever classify worker asked first - up to eight racing into one process-wide cache. The load path now warms the backbone single-threaded under the existing lock, on the cache-hit branch too. The readiness path loads the gallery before the socket accepts, so the warmup is free. It is Metal-only and best-effort: a CPU gallery has no shader library to compile, and a cache hit is served without a usable backbone by contract. The regression test records overlap rather than timing: a fake backbone that sleeps and counts concurrent entries, run from four threads, asserting the count never exceeds one. It fails on the previous gallery.py and needs no MPS.
631 lines
27 KiB
Python
631 lines
27 KiB
Python
"""
|
|
Tests for the product gallery's matching logic.
|
|
|
|
The embedding model is injected, so these run without torch or a GPU: what is
|
|
being tested is the retrieval contract - best view wins, several views per SKU
|
|
count as one answer, and a stranger is rejected rather than named.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import os
|
|
import sys
|
|
import unittest
|
|
|
|
import numpy as np
|
|
|
|
current_dir = os.path.dirname(os.path.abspath(__file__))
|
|
project_root = os.path.dirname(current_dir)
|
|
if os.path.exists(os.path.join(project_root, "stock")):
|
|
os.chdir(project_root)
|
|
if current_dir not in sys.path:
|
|
sys.path.insert(0, current_dir)
|
|
|
|
from gallery import ( # noqa: E402
|
|
DINOV2_MODEL,
|
|
DINOV2_SHA256,
|
|
GalleryEntry,
|
|
ProductGallery,
|
|
catalog_digest,
|
|
dinov2_weights_path,
|
|
head_class_set,
|
|
load_dinov2,
|
|
)
|
|
|
|
import gallery as gallery_module # noqa: E402
|
|
|
|
|
|
def a_gallery(views, threshold=0.55):
|
|
"""A gallery with hand-made embeddings: no model, no files."""
|
|
skus = sorted({sku for sku, _ in views})
|
|
gallery = ProductGallery("does-not-exist", skus, threshold=threshold)
|
|
embeddings = []
|
|
entries = []
|
|
for index, (sku, vector) in enumerate(views):
|
|
vector = np.asarray(vector, dtype=np.float32)
|
|
vector /= np.linalg.norm(vector)
|
|
embeddings.append(vector)
|
|
entries.append(GalleryEntry(sku_id=sku, view=f"{sku}_{index}", embedding=vector))
|
|
gallery._entries = entries
|
|
gallery._matrix = np.stack(embeddings)
|
|
gallery._available = True
|
|
return gallery
|
|
|
|
|
|
def unit(vector):
|
|
vector = np.asarray(vector, dtype=np.float32)
|
|
return (vector / np.linalg.norm(vector)).tolist()
|
|
|
|
|
|
class Backbone:
|
|
"""A deterministic stand-in for DINOv2: the prepared crop's channel means,
|
|
tiled to the width the real backbone returns."""
|
|
|
|
def __init__(self):
|
|
self.calls = 0
|
|
|
|
def __call__(self, tensor):
|
|
self.calls += 1
|
|
return tensor.mean(dim=(2, 3)).repeat(1, gallery_module.DINOV2_EMBED_DIM // 3)
|
|
|
|
|
|
def a_query():
|
|
"""A unit vector in the backbone's space, for match() without a model."""
|
|
query = np.zeros(gallery_module.DINOV2_EMBED_DIM, dtype=np.float32)
|
|
query[0] = 1.0
|
|
query[1] = 0.05
|
|
return query / np.linalg.norm(query)
|
|
|
|
|
|
class UnreachableBackbone:
|
|
def __call__(self, tensor):
|
|
raise AssertionError("the gallery loaded a backbone it should have cached")
|
|
|
|
|
|
def write_view(path, seed):
|
|
import cv2
|
|
|
|
image = np.random.RandomState(seed).randint(0, 256, (24, 32, 3), dtype=np.uint8)
|
|
cv2.imwrite(path, image)
|
|
|
|
|
|
def a_store(directory, names=("sku_a", "sku_a_1", "sku_b")):
|
|
"""A store dir: one JPEG per view plus the stock.csv the catalog is read from."""
|
|
with open(os.path.join(directory, "stock.csv"), "w", encoding="utf-8") as handle:
|
|
handle.write("SKU_ID,product_name\n")
|
|
for index, name in enumerate(names):
|
|
write_view(os.path.join(directory, f"{name}.jpeg"), index)
|
|
|
|
|
|
def a_stored_gallery(directory, names, backbone, image_size=32):
|
|
gallery = ProductGallery(directory, list(names), device="cpu", image_size=image_size)
|
|
gallery._load_model = lambda: backbone
|
|
return gallery
|
|
|
|
|
|
class MatchingTests(unittest.TestCase):
|
|
def test_the_nearest_view_names_the_product(self):
|
|
gallery = a_gallery([("sku_a", unit([1, 0, 0])), ("sku_b", unit([0, 1, 0]))])
|
|
gallery.embed = lambda crop: np.array([0.9, 0.1, 0.0], dtype=np.float32)
|
|
match = gallery.match(object())
|
|
self.assertFalse(match.rejected)
|
|
self.assertEqual(match.sku_id, "sku_a")
|
|
self.assertAlmostEqual(match.similarity, 0.9, places=3)
|
|
|
|
def test_several_views_of_one_sku_are_one_answer(self):
|
|
# The second view of sku_b is closer than any view of sku_a, but sku_a must
|
|
# still win on its best view: a SKU is not outvoted by having more photos.
|
|
gallery = a_gallery([
|
|
("sku_a", unit([1, 0, 0])),
|
|
("sku_b", unit([0, 1, 0])),
|
|
("sku_b", unit([0.2, 0.98, 0])),
|
|
])
|
|
gallery.embed = lambda crop: np.array([1.0, 0.05, 0.0], dtype=np.float32)
|
|
match = gallery.match(object())
|
|
self.assertEqual(match.sku_id, "sku_a")
|
|
self.assertEqual(match.runner_up, "sku_b")
|
|
|
|
def test_a_stranger_is_rejected_not_named(self):
|
|
gallery = a_gallery([("sku_a", unit([1, 0, 0])), ("sku_b", unit([0, 1, 0]))])
|
|
gallery.embed = lambda crop: np.array([0.0, 0.0, 1.0], dtype=np.float32)
|
|
match = gallery.match(object())
|
|
self.assertTrue(match.rejected)
|
|
self.assertIsNone(match.sku_id)
|
|
|
|
def test_the_threshold_moves_the_rejection_line(self):
|
|
gallery = a_gallery([("sku_a", unit([1, 0, 0]))], threshold=0.99)
|
|
gallery.embed = lambda crop: np.array([0.9, 0.4, 0.0], dtype=np.float32)
|
|
self.assertTrue(gallery.match(object()).rejected)
|
|
self.assertFalse(gallery.match(object(), threshold=0.5).rejected)
|
|
|
|
def test_margin_separates_a_confident_match_from_a_close_one(self):
|
|
gallery = a_gallery([("sku_a", unit([1, 0, 0])), ("sku_b", unit([0.9, 0.44, 0]))])
|
|
gallery.embed = lambda crop: np.array([1.0, 0.1, 0.0], dtype=np.float32)
|
|
confident = gallery.match(object())
|
|
gallery.embed = lambda crop: np.array([0.95, 0.3, 0.0], dtype=np.float32)
|
|
close = gallery.match(object())
|
|
self.assertGreater(confident.margin, close.margin)
|
|
|
|
def test_the_per_sku_max_does_not_depend_on_entry_order(self):
|
|
# Three SKUs, several views each, enrolled out of order: a SKU is scored
|
|
# by its best view, and the answer cannot move because a row moved.
|
|
gallery = a_gallery([
|
|
("sku_b", unit([0, 1, 0])),
|
|
("sku_a", unit([1, 0, 0])),
|
|
("sku_c", unit([0, 0, 1])),
|
|
("sku_b", unit([0.3, 0.95, 0])),
|
|
("sku_a", unit([0.98, 0.2, 0])),
|
|
])
|
|
gallery.embed = lambda crop: np.array([1.0, 0.05, 0.0], dtype=np.float32)
|
|
straight = gallery.match(object())
|
|
self.assertEqual(straight.sku_id, "sku_a")
|
|
self.assertEqual(straight.runner_up, "sku_b")
|
|
|
|
for order in ([4, 0, 3, 1, 2], [2, 3, 1, 4, 0]):
|
|
gallery._entries = [gallery._entries[index] for index in order]
|
|
gallery._matrix = gallery._matrix[order]
|
|
shuffled = gallery.match(object())
|
|
self.assertEqual(shuffled.sku_id, straight.sku_id)
|
|
self.assertEqual(shuffled.runner_up, straight.runner_up)
|
|
self.assertAlmostEqual(shuffled.similarity, straight.similarity, places=6)
|
|
self.assertAlmostEqual(shuffled.runner_up_similarity, straight.runner_up_similarity, places=6)
|
|
|
|
def test_a_tie_keeps_the_first_enrolled_sku(self):
|
|
# Exactly equidistant: first appearance decides, as it did before the
|
|
# per-SKU maximum became a segmented reduction.
|
|
gallery = a_gallery([("sku_a", unit([1, 0, 0])), ("sku_b", unit([0, 1, 0]))])
|
|
gallery.embed = lambda crop: np.array([1.0, 1.0, 0.0], dtype=np.float32)
|
|
match = gallery.match(object())
|
|
self.assertFalse(match.rejected)
|
|
self.assertEqual(match.sku_id, "sku_a")
|
|
self.assertEqual(match.runner_up, "sku_b")
|
|
self.assertAlmostEqual(match.similarity, match.runner_up_similarity, places=6)
|
|
self.assertAlmostEqual(match.margin, 0.0, places=6)
|
|
|
|
def test_an_unloaded_gallery_rejects_everything(self):
|
|
gallery = ProductGallery("does-not-exist", ["sku_a"])
|
|
match = gallery.match(object())
|
|
self.assertTrue(match.rejected)
|
|
|
|
|
|
class EnrollmentTests(unittest.TestCase):
|
|
def test_views_are_found_by_sku_prefix(self):
|
|
import tempfile
|
|
|
|
with tempfile.TemporaryDirectory() as directory:
|
|
for name in ("sku_a.jpeg", "sku_a_1.jpeg", "sku_a-2.jpeg", "sku_ab.jpeg", "notes.txt"):
|
|
open(os.path.join(directory, name), "wb").close()
|
|
gallery = ProductGallery(directory, ["sku_a", "sku_ab"])
|
|
self.assertEqual(len(gallery._views_for("sku_a")), 3,
|
|
"one view, two extra views, and not sku_ab")
|
|
self.assertEqual(len(gallery._views_for("sku_ab")), 1)
|
|
|
|
|
|
class FeatureCacheTests(unittest.TestCase):
|
|
"""The cached matrix must be this store's, or it must not be used."""
|
|
|
|
def setUp(self):
|
|
import hashlib
|
|
import tempfile
|
|
|
|
self.directory = tempfile.TemporaryDirectory()
|
|
self.addCleanup(self.directory.cleanup)
|
|
a_store(self.directory.name)
|
|
# A cache is only served when the pinned backbone is on disk, so point the
|
|
# pin at a stand-in file the module's expected digest is told about. The
|
|
# check itself - hash this file, compare - is what is under test.
|
|
self.weights = os.path.join(self.directory.name, "dinov2.safetensors")
|
|
self.write(self.weights, b"not really a checkpoint")
|
|
os.environ["DORAYAKI_DINOV2_WEIGHTS"] = self.weights
|
|
self.addCleanup(os.environ.pop, "DORAYAKI_DINOV2_WEIGHTS", None)
|
|
self.addCleanup(setattr, gallery_module, "DINOV2_SHA256", gallery_module.DINOV2_SHA256)
|
|
gallery_module.DINOV2_SHA256 = hashlib.sha256(b"not really a checkpoint").hexdigest()
|
|
|
|
def write(self, path, content):
|
|
with open(path, "wb") as handle:
|
|
handle.write(content)
|
|
return path
|
|
|
|
def real_gallery(self):
|
|
"""A gallery with the real _load_model: the server's own readiness path."""
|
|
return ProductGallery(self.directory.name, ["sku_a", "sku_b"], device="cpu", image_size=32)
|
|
|
|
def loaded(self, names=("sku_a", "sku_b"), image_size=32):
|
|
backbone = Backbone()
|
|
gallery = a_stored_gallery(self.directory.name, names, backbone, image_size=image_size)
|
|
self.assertTrue(gallery.load())
|
|
return gallery, backbone
|
|
|
|
def cache_payload(self, gallery):
|
|
"""The rows and matrix a loaded gallery left on disk."""
|
|
with np.load(gallery._cache_path(), allow_pickle=False) as payload:
|
|
return (
|
|
[str(value) for value in payload["skus"].tolist()],
|
|
[str(value) for value in payload["views"].tolist()],
|
|
payload["matrix"],
|
|
)
|
|
|
|
def rewrite_cache(self, gallery, skus, views, matrix):
|
|
"""The same valid key with different contents: only the rows can catch it."""
|
|
path = gallery._cache_path()
|
|
with np.load(path, allow_pickle=False) as payload:
|
|
key = str(payload["key"].item())
|
|
np.savez(path, key=key, skus=np.array(skus), views=np.array(views), matrix=matrix)
|
|
|
|
def re_embeds(self, first, why):
|
|
"""The one right answer to a rewritten cache: build it from the pixels."""
|
|
backbone = Backbone()
|
|
gallery = a_stored_gallery(self.directory.name, ("sku_a", "sku_b"), backbone)
|
|
self.assertTrue(gallery.load(), why)
|
|
self.assertEqual(backbone.calls, 1, f"{why} must not be served from cache")
|
|
np.testing.assert_array_equal(gallery._matrix, first._matrix)
|
|
self.assertEqual([entry.sku_id for entry in gallery._entries], ["sku_a", "sku_a", "sku_b"])
|
|
|
|
def test_a_second_start_answers_the_same_from_the_cache(self):
|
|
first, backbone = self.loaded()
|
|
self.assertEqual(backbone.calls, 1)
|
|
self.assertEqual(first.view_count(), 3)
|
|
self.assertTrue(os.path.isfile(first._cache_path()))
|
|
|
|
cached = a_stored_gallery(self.directory.name, ("sku_a", "sku_b"), UnreachableBackbone())
|
|
self.assertTrue(cached.load())
|
|
self.assertEqual(cached.view_count(), first.view_count())
|
|
np.testing.assert_array_equal(cached._matrix, first._matrix)
|
|
self.assertEqual([entry.view for entry in cached._entries],
|
|
[entry.view for entry in first._entries])
|
|
query = a_query()
|
|
first.embed = cached.embed = lambda crop: query
|
|
self.assertEqual(vars(cached.match(object())), vars(first.match(object())))
|
|
|
|
def test_changed_view_bytes_re_embed(self):
|
|
first, _ = self.loaded()
|
|
before = first._matrix.copy()
|
|
write_view(os.path.join(self.directory.name, "sku_b.jpeg"), seed=99)
|
|
|
|
second, backbone = self.loaded()
|
|
self.assertEqual(backbone.calls, 1, "a replaced view must not be served from cache")
|
|
self.assertFalse(np.array_equal(second._matrix, before))
|
|
|
|
def test_a_reordered_catalog_does_not_borrow_another_mapping(self):
|
|
self.loaded()
|
|
swapped, backbone = self.loaded(names=("sku_b", "sku_a"))
|
|
self.assertEqual(backbone.calls, 1)
|
|
self.assertEqual([entry.sku_id for entry in swapped._entries], ["sku_b", "sku_a", "sku_a"])
|
|
|
|
def test_another_image_size_is_never_served(self):
|
|
small, _ = self.loaded(image_size=32)
|
|
large, backbone = self.loaded(image_size=64)
|
|
self.assertEqual(backbone.calls, 1)
|
|
self.assertNotEqual(small._cache_path(), large._cache_path())
|
|
|
|
def test_a_corrupt_cache_re_embeds_the_same_matrix(self):
|
|
first, _ = self.loaded()
|
|
with open(first._cache_path(), "wb") as handle:
|
|
handle.write(b"not an npz at all")
|
|
|
|
second, backbone = self.loaded()
|
|
self.assertEqual(backbone.calls, 1)
|
|
np.testing.assert_array_equal(second._matrix, first._matrix)
|
|
|
|
def test_swapped_cache_rows_are_not_served(self):
|
|
# A current key says nothing about the rows inside the file. The two
|
|
# sku_a views are adjacent, so a swap has to cross SKUs to mean
|
|
# anything - and the matrix that moves with them is the dangerous part:
|
|
# a vector that would then be filed under the wrong product.
|
|
first, _ = self.loaded()
|
|
skus, views, matrix = self.cache_payload(first)
|
|
self.assertEqual(list(zip(skus, views)),
|
|
[("sku_a", "sku_a.jpeg"), ("sku_a", "sku_a_1.jpeg"), ("sku_b", "sku_b.jpeg")])
|
|
self.rewrite_cache(first, skus[::-1], views[::-1], matrix[::-1])
|
|
self.re_embeds(first, "rows swapped with their vectors")
|
|
|
|
def test_mislabelled_cache_rows_are_not_served(self):
|
|
# The views and the matrix stay where they are, so every row names a
|
|
# product it was not embedded for.
|
|
first, _ = self.loaded()
|
|
skus, views, matrix = self.cache_payload(first)
|
|
self.rewrite_cache(first, [skus[2], skus[1], skus[0]], views, matrix)
|
|
self.re_embeds(first, "labels swapped around the views")
|
|
|
|
def test_a_duplicated_cache_row_is_not_served(self):
|
|
# The first view written twice: one image would answer for two products.
|
|
first, _ = self.loaded()
|
|
skus, views, matrix = self.cache_payload(first)
|
|
self.rewrite_cache(first, skus[:1] * 2 + skus[2:], views[:1] * 2 + views[2:], matrix)
|
|
self.re_embeds(first, "a row repeated")
|
|
|
|
def test_a_foreign_sku_row_is_not_served(self):
|
|
# A row from another store's catalog: a SKU this gallery never enrolled.
|
|
first, _ = self.loaded()
|
|
skus, views, matrix = self.cache_payload(first)
|
|
self.rewrite_cache(first, skus[1:] + ["other_store"], views[1:] + ["other_store.jpeg"], matrix)
|
|
self.re_embeds(first, "another store's row")
|
|
|
|
def test_a_cache_may_omit_a_view_it_could_not_decode(self):
|
|
# The writer drops a view it could not read and keeps everyone else's
|
|
# order, so such a cache is still this store's - and still worth serving
|
|
# rather than re-embedding the whole catalog on every start.
|
|
first, _ = self.loaded()
|
|
skus, views, matrix = self.cache_payload(first)
|
|
self.rewrite_cache(first, skus[:1] + skus[2:], views[:1] + views[2:], matrix[[0, 2]])
|
|
|
|
backbone = Backbone()
|
|
gallery = a_stored_gallery(self.directory.name, ("sku_a", "sku_b"), backbone)
|
|
self.assertTrue(gallery.load())
|
|
self.assertEqual(backbone.calls, 0, "a readable shorter cache is still a hit")
|
|
self.assertEqual([entry.sku_id for entry in gallery._entries], ["sku_a", "sku_b"])
|
|
np.testing.assert_array_equal(gallery._matrix, first._matrix[[0, 2]])
|
|
|
|
def test_a_nonnumeric_cache_matrix_is_not_served(self):
|
|
# Only the float matrix _cache_write produces is a cache. A complex
|
|
# matrix is finite and casts without complaint, so it answers with a
|
|
# different matrix than the one on disk; integer, boolean and text
|
|
# matrices are not embeddings at all.
|
|
first, _ = self.loaded()
|
|
skus, views, matrix = self.cache_payload(first)
|
|
others = {
|
|
"complex": matrix.astype(np.complex128),
|
|
"timedelta": np.zeros_like(matrix, dtype="timedelta64[s]"),
|
|
"integer": (matrix * 8).astype(np.int32),
|
|
"boolean": np.ones_like(matrix, dtype=bool),
|
|
"text": np.zeros(matrix.shape, dtype="<U4"),
|
|
}
|
|
for name, other in others.items():
|
|
with self.subTest(matrix=name):
|
|
self.rewrite_cache(first, skus, views, other)
|
|
self.re_embeds(first, f"a {name} matrix")
|
|
|
|
def test_a_loaded_gallery_reloads_when_its_source_moves(self):
|
|
gallery, backbone = self.loaded()
|
|
self.assertTrue(gallery.load())
|
|
self.assertEqual(backbone.calls, 1, "an unchanged store must not re-embed")
|
|
|
|
write_view(os.path.join(self.directory.name, "sku_b.jpeg"), seed=99)
|
|
self.assertTrue(gallery.load())
|
|
self.assertEqual(backbone.calls, 2)
|
|
|
|
with open(os.path.join(self.directory.name, "stock.csv"), "a", encoding="utf-8") as handle:
|
|
handle.write("sku_c,new product\n")
|
|
self.assertTrue(gallery.load())
|
|
self.assertEqual(backbone.calls, 3, "stock.csv is part of what an embedding is pinned to")
|
|
self.assertEqual(gallery.view_count(), 3)
|
|
|
|
def test_a_valid_cache_needs_the_pinned_backbone_on_disk(self):
|
|
# The regression: a cache hit never opens load_dinov2, so a deleted or
|
|
# altered checkpoint used to read as a loaded gallery and then raise on
|
|
# the first match - after the server had reported itself ready.
|
|
self.loaded()
|
|
gallery = self.real_gallery()
|
|
self.assertTrue(gallery.load(), "the cache serves while the pin holds")
|
|
|
|
os.environ["DORAYAKI_DINOV2_WEIGHTS"] = os.path.join(self.directory.name, "gone.safetensors")
|
|
missing = self.real_gallery()
|
|
self.assertFalse(missing.load())
|
|
self.assertFalse(missing.available())
|
|
|
|
os.environ["DORAYAKI_DINOV2_WEIGHTS"] = self.write(
|
|
os.path.join(self.directory.name, "altered.safetensors"), b"a different checkpoint"
|
|
)
|
|
altered = self.real_gallery()
|
|
self.assertFalse(altered.load())
|
|
self.assertFalse(altered.available())
|
|
|
|
def test_an_altered_backbone_re_embeds_rather_than_serves(self):
|
|
# With a backbone that can still load, an altered file must re-embed.
|
|
self.loaded()
|
|
self.assertTrue(
|
|
a_stored_gallery(self.directory.name, ("sku_a", "sku_b"), UnreachableBackbone()).load(),
|
|
"the cache still serves on the pinned file",
|
|
)
|
|
self.write(self.weights, b"a different checkpoint")
|
|
backbone = Backbone()
|
|
rebuilt = a_stored_gallery(self.directory.name, ("sku_a", "sku_b"), backbone)
|
|
self.assertTrue(rebuilt.load())
|
|
self.assertEqual(backbone.calls, 1, "an unpinned cache must not be served")
|
|
|
|
|
|
class OfflineWeightsTests(unittest.TestCase):
|
|
def test_a_missing_weights_file_names_the_path_and_does_not_fetch(self):
|
|
previous = os.environ.get("DORAYAKI_DINOV2_WEIGHTS")
|
|
os.environ["DORAYAKI_DINOV2_WEIGHTS"] = os.path.join(
|
|
"does-not-exist", "dinov2_vits14_lvd142m.safetensors"
|
|
)
|
|
try:
|
|
with self.assertRaises(FileNotFoundError) as raised:
|
|
load_dinov2(224)
|
|
finally:
|
|
if previous is None:
|
|
os.environ.pop("DORAYAKI_DINOV2_WEIGHTS", None)
|
|
else:
|
|
os.environ["DORAYAKI_DINOV2_WEIGHTS"] = previous
|
|
message = str(raised.exception)
|
|
self.assertIn("does-not-exist", message)
|
|
self.assertIn(DINOV2_SHA256, message)
|
|
self.assertIn("does not download", message)
|
|
|
|
def test_the_default_path_stays_inside_the_repo(self):
|
|
previous = os.environ.pop("DORAYAKI_DINOV2_WEIGHTS", None)
|
|
try:
|
|
self.assertTrue(dinov2_weights_path().endswith(
|
|
os.path.join("weights", "dinov2_vits14_lvd142m.safetensors")
|
|
))
|
|
finally:
|
|
if previous is not None:
|
|
os.environ["DORAYAKI_DINOV2_WEIGHTS"] = previous
|
|
|
|
|
|
def _head_payload(class_ids, torch):
|
|
projection = torch.nn.Linear(384, 128)
|
|
prototypes = torch.nn.functional.normalize(torch.randn(len(class_ids), 128), dim=1)
|
|
return {
|
|
"class_ids": list(class_ids),
|
|
"image_size": 518,
|
|
"base_model": DINOV2_MODEL,
|
|
"projection": {key: value.detach().cpu() for key, value in projection.state_dict().items()},
|
|
"prototypes": prototypes.detach().cpu(),
|
|
}
|
|
|
|
|
|
class HeadPinTests(unittest.TestCase):
|
|
def test_a_stale_head_cannot_label_this_catalog(self):
|
|
import tempfile
|
|
|
|
import torch
|
|
|
|
with tempfile.TemporaryDirectory() as directory:
|
|
with open(os.path.join(directory, "stock.csv"), "w", encoding="utf-8") as handle:
|
|
handle.write("SKU_ID,product_name\nsku_a,A\nsku_b,B\nsku_c,C\n")
|
|
gallery = ProductGallery(directory, ["sku_a", "sku_b", "sku_c"])
|
|
gallery._available = True
|
|
path = gallery._head_path
|
|
subset = ["sku_c", "sku_a"]
|
|
digest = catalog_digest(directory, subset)
|
|
|
|
foreign = _head_payload(["other_store", "sku_a"], torch)
|
|
torch.save(foreign, path)
|
|
self.assertFalse(gallery.load_head())
|
|
self.assertIsNone(gallery._head)
|
|
self.assertIn("catalog_mismatch", gallery.head_reason)
|
|
self.assertTrue(gallery.available())
|
|
|
|
legacy = _head_payload(subset, torch)
|
|
torch.save(legacy, path)
|
|
self.assertTrue(gallery.load_head())
|
|
self.assertEqual(gallery.head_reason, "legacy_unpinned")
|
|
self.assertEqual(gallery._head["class_ids"], subset)
|
|
self.assertNotIn("other_store", gallery._head["class_ids"])
|
|
|
|
stale = _head_payload(subset, torch)
|
|
stale["catalog_digest"] = "0" * 64
|
|
torch.save(stale, path)
|
|
self.assertFalse(gallery.load_head())
|
|
self.assertIsNone(gallery._head)
|
|
self.assertIn("catalog_mismatch", gallery.head_reason)
|
|
self.assertTrue(gallery.available())
|
|
|
|
wrong_backbone = _head_payload(subset, torch)
|
|
wrong_backbone["catalog_digest"] = digest
|
|
wrong_backbone["base_model"] = "vit_base_patch14_dinov2.lvd142m"
|
|
torch.save(wrong_backbone, path)
|
|
self.assertFalse(gallery.load_head())
|
|
self.assertIn("backbone_mismatch", gallery.head_reason)
|
|
|
|
wrong_size = _head_payload(subset, torch)
|
|
wrong_size["catalog_digest"] = digest
|
|
wrong_size["image_size"] = 224
|
|
torch.save(wrong_size, path)
|
|
self.assertFalse(gallery.load_head())
|
|
self.assertIn("image_size_mismatch", gallery.head_reason)
|
|
|
|
torch.save(["not", "a", "head"], path)
|
|
self.assertFalse(gallery.load_head())
|
|
self.assertTrue(gallery.head_reason.startswith("malformed"))
|
|
self.assertTrue(gallery.available())
|
|
self.assertIsNone(gallery._head)
|
|
|
|
current = _head_payload(subset, torch)
|
|
current["catalog_digest"] = digest
|
|
torch.save(current, path)
|
|
self.assertTrue(gallery.load_head())
|
|
self.assertEqual(gallery.head_reason, "ok")
|
|
self.assertEqual(gallery._head["class_ids"], subset)
|
|
self.assertIsNone(head_class_set(["sku_a", "sku_a", ""]))
|
|
|
|
|
|
|
|
class ConcurrentInferenceTests(unittest.TestCase):
|
|
"""The backbone is never entered by two threads at once.
|
|
|
|
Metal's shader library is process-global and is not thread-safe while
|
|
kernels are still being compiled into it, so two threads reaching the first
|
|
forward together segfault the whole server. This is a contract test, not a
|
|
timing test: it records overlap rather than measuring speed.
|
|
"""
|
|
|
|
def test_two_threads_never_enter_the_backbone_together(self):
|
|
import time
|
|
import threading
|
|
import tempfile
|
|
|
|
with tempfile.TemporaryDirectory() as directory:
|
|
a_store(directory, names=("sku_a", "sku_b"))
|
|
gallery = ProductGallery(directory, ["sku_a", "sku_b"], device="cpu", image_size=32)
|
|
|
|
inside = 0
|
|
overlapped = []
|
|
gate = threading.Lock()
|
|
|
|
class SlowBackbone:
|
|
def __call__(self, tensor):
|
|
nonlocal inside
|
|
with gate:
|
|
inside += 1
|
|
if inside > 1:
|
|
overlapped.append(inside)
|
|
time.sleep(0.01)
|
|
with gate:
|
|
inside -= 1
|
|
return tensor.mean(dim=(2, 3)).repeat(1, gallery_module.DINOV2_EMBED_DIM // 3)
|
|
|
|
gallery._load_model = lambda: SlowBackbone()
|
|
crop = np.zeros((24, 32, 3), dtype=np.uint8)
|
|
errors = []
|
|
|
|
def work():
|
|
try:
|
|
for _ in range(6):
|
|
gallery.embed(crop)
|
|
except Exception as error: # pragma: no cover - reported below
|
|
errors.append(repr(error))
|
|
|
|
threads = [threading.Thread(target=work) for _ in range(4)]
|
|
for thread in threads:
|
|
thread.start()
|
|
for thread in threads:
|
|
thread.join()
|
|
|
|
self.assertEqual(errors, [])
|
|
self.assertEqual(overlapped, [], "two threads ran the backbone at the same time")
|
|
|
|
|
|
class TrainerPinTests(unittest.TestCase):
|
|
def test_head_training_loads_the_pinned_file_and_keeps_preprocess(self):
|
|
import finetune_identity
|
|
import gallery as gallery_module
|
|
|
|
seen = []
|
|
|
|
class FakeBackbone:
|
|
def eval(self):
|
|
return self
|
|
|
|
def to(self, device):
|
|
return self
|
|
|
|
sentinel = FakeBackbone()
|
|
def fake_load(image_size):
|
|
seen.append(image_size)
|
|
return sentinel
|
|
|
|
|
|
# build_model resolves the loader from gallery at call time, so that is
|
|
# the attribute to replace; finetune_identity has no module-level copy.
|
|
original = gallery_module.load_dinov2
|
|
gallery_module.load_dinov2 = fake_load
|
|
try:
|
|
model, device = finetune_identity.build_model(224)
|
|
finally:
|
|
gallery_module.load_dinov2 = original
|
|
self.assertIs(model, sentinel)
|
|
self.assertIn(device, ("cpu", "mps"))
|
|
self.assertEqual(seen, [224])
|
|
|
|
image = np.zeros((10, 20, 3), dtype=np.uint8)
|
|
self.assertEqual(finetune_identity.prepare(image, 32).shape, gallery_module._prepare(image, 32).shape)
|
|
self.assertTrue(np.allclose(
|
|
finetune_identity.prepare(image, 32),
|
|
gallery_module._prepare(image, 32),
|
|
))
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|