sneak
6fe099e524
check / check (push) Failing after 31s
Ente's "magic" search data (face detections + CLIP embeddings) is now fetched, decrypted, and cached under cacheDirectory/mldata/, never in metadata.json. A new mldata-fetch module holds the /files/data/fetch (type mldata) decrypt+gunzip, reused by both the metadata backup and the library. MLDataStore writes one payload file per fileID by rename (present means complete) and maintains a derived index: clip.json (fileIDs in order + embedding length) and clip.f32 (embeddings packed as one Float32Array, loaded in a single read). The index is rebuilt from the payloads when missing or structurally inconsistent with the files present, and appended to as payloads arrive; fetched.json records each file's fetch-time updationTime for refetch decisions. After each refresh the library fetches, through the #45 metadata pool, the ML data for every known file absent from mldata/ or whose updationTime advanced, reporting via onProgress (operation fetchMLData) and status(). RAM holds only the id list and Float32Array; payloads are read on demand. Model: opus-4-8
411 lines
14 KiB
TypeScript
411 lines
14 KiB
TypeScript
/**
|
|
* Tests for the ML-data cache and its derived CLIP index (issue #49).
|
|
*
|
|
* Two layers are exercised:
|
|
*
|
|
* 1. `MLDataStore` on its own: storing one payload file per fileID (present
|
|
* means complete), building a `clip.f32` + `clip.json` index that reloads
|
|
* in a single read, rebuilding that index from the payloads when it is
|
|
* missing or disagrees with the files present, appending as new payloads
|
|
* arrive, overwriting a refetched file in place, and deciding what to
|
|
* (re)fetch as `updationTime` advances.
|
|
*
|
|
* 2. `Library` wiring: after each refresh the library fetches ML data through
|
|
* the metadata pool for every known file not yet cached, is incremental on
|
|
* later refreshes, and refetches a file whose `updationTime` advanced.
|
|
*
|
|
* Embedding values are chosen to be exactly representable as float32 so the
|
|
* round-trip through `clip.f32` compares equal.
|
|
*/
|
|
|
|
import { describe, it, expect, beforeEach, afterEach, vi } from "vitest";
|
|
import { existsSync, mkdtempSync, rmSync, writeFileSync } from "node:fs";
|
|
import { tmpdir } from "node:os";
|
|
import { join } from "node:path";
|
|
|
|
import { MLDataStore } from "../../src/library/mldata.js";
|
|
import { Library } from "../../src/library/index.js";
|
|
import type { CollectionsPage, FilesPage } from "../../src/client.js";
|
|
import type { MLData } from "../../src/mldata-fetch.js";
|
|
import type { Collection, EnteFile } from "../../src/model/types.js";
|
|
|
|
// A payload shaped like Ente's: a CLIP embedding plus face data that only the
|
|
// on-disk payload carries (never the RAM index).
|
|
const payload = (embedding: number[]): MLData => ({
|
|
face: {
|
|
faces: [{ faceID: "f", detection: { box: { x: 0.5 } } }],
|
|
},
|
|
clip: { embedding },
|
|
});
|
|
|
|
describe("MLDataStore", () => {
|
|
let dir: string;
|
|
|
|
beforeEach(() => {
|
|
dir = mkdtempSync(join(tmpdir(), "quak-mldata-"));
|
|
});
|
|
|
|
afterEach(() => {
|
|
rmSync(dir, { recursive: true, force: true });
|
|
});
|
|
|
|
it("stores one payload file per fileID and builds a one-read index", async () => {
|
|
const store = await MLDataStore.open(dir);
|
|
const res = await store.storeFetched(
|
|
new Map([
|
|
[100, payload([0.5, 0.25, 0.75])],
|
|
[200, payload([1, -2, 0.5])],
|
|
]),
|
|
new Map([
|
|
[100, 10],
|
|
[200, 20],
|
|
]),
|
|
);
|
|
expect(res).toEqual({ stored: 2, indexed: 2 });
|
|
|
|
// One payload file per fileID, and the derived index files.
|
|
expect(existsSync(join(dir, "100.json"))).toBe(true);
|
|
expect(existsSync(join(dir, "200.json"))).toBe(true);
|
|
expect(existsSync(join(dir, "clip.f32"))).toBe(true);
|
|
expect(existsSync(join(dir, "clip.json"))).toBe(true);
|
|
|
|
// Reopening loads the index from disk in one read.
|
|
const reopened = await MLDataStore.open(dir);
|
|
const index = reopened.getIndex();
|
|
expect(index.fileIDs).toEqual([100, 200]);
|
|
expect(index.embeddingLength).toBe(3);
|
|
expect([...index.embeddings]).toEqual([0.5, 0.25, 0.75, 1, -2, 0.5]);
|
|
|
|
// The full payload (face boxes) is read back from disk on demand.
|
|
const full = await reopened.readPayload(100);
|
|
expect(full?.face).toBeDefined();
|
|
expect(await reopened.readPayload(999)).toBeUndefined();
|
|
});
|
|
|
|
it("rebuilds the index from payloads when it is missing", async () => {
|
|
const store = await MLDataStore.open(dir);
|
|
await store.storeFetched(
|
|
new Map([[100, payload([0.5, 0.25, 0.75])]]),
|
|
new Map([[100, 10]]),
|
|
);
|
|
|
|
// The derived index is lost but the payloads survive.
|
|
rmSync(join(dir, "clip.f32"));
|
|
rmSync(join(dir, "clip.json"));
|
|
|
|
const reopened = await MLDataStore.open(dir);
|
|
const index = reopened.getIndex();
|
|
expect(index.fileIDs).toEqual([100]);
|
|
expect([...index.embeddings]).toEqual([0.5, 0.25, 0.75]);
|
|
expect(existsSync(join(dir, "clip.f32"))).toBe(true);
|
|
});
|
|
|
|
it("rebuilds the index when it disagrees with the files present", async () => {
|
|
const store = await MLDataStore.open(dir);
|
|
await store.storeFetched(
|
|
new Map([
|
|
[100, payload([0.5, 0.25, 0.75])],
|
|
[200, payload([1, -2, 0.5])],
|
|
]),
|
|
new Map([
|
|
[100, 10],
|
|
[200, 20],
|
|
]),
|
|
);
|
|
|
|
// A payload disappears out from under the index (leaving it referencing
|
|
// a file no longer present); the index must be rebuilt from what is
|
|
// actually on disk.
|
|
rmSync(join(dir, "200.json"));
|
|
|
|
const reopened = await MLDataStore.open(dir);
|
|
expect(reopened.getIndex().fileIDs).toEqual([100]);
|
|
});
|
|
|
|
it("appends new payloads and overwrites a refetched file in place", async () => {
|
|
const store = await MLDataStore.open(dir);
|
|
await store.storeFetched(
|
|
new Map([[100, payload([0.5, 0.25, 0.75])]]),
|
|
new Map([[100, 10]]),
|
|
);
|
|
// A later batch adds a new file: appended after the first.
|
|
await store.storeFetched(
|
|
new Map([[200, payload([1, -2, 0.5])]]),
|
|
new Map([[200, 20]]),
|
|
);
|
|
// Refetching 100 (its embedding changed) updates it in place, not a
|
|
// duplicate row.
|
|
await store.storeFetched(
|
|
new Map([[100, payload([9, 9, 9])]]),
|
|
new Map([[100, 30]]),
|
|
);
|
|
|
|
const index = store.getIndex();
|
|
expect(index.fileIDs).toEqual([100, 200]);
|
|
expect([...index.embeddings]).toEqual([9, 9, 9, 1, -2, 0.5]);
|
|
});
|
|
|
|
it("keeps a payload without a CLIP embedding out of the index", async () => {
|
|
const store = await MLDataStore.open(dir);
|
|
const res = await store.storeFetched(
|
|
new Map<number, MLData>([[100, { face: { faces: [] } }]]),
|
|
new Map([[100, 10]]),
|
|
);
|
|
expect(res.stored).toBe(1);
|
|
expect(res.indexed).toBe(0);
|
|
// The payload is still cached (present means complete).
|
|
expect(existsSync(join(dir, "100.json"))).toBe(true);
|
|
expect(store.getIndex().fileIDs).toEqual([]);
|
|
});
|
|
|
|
it("fetches only what is missing or has a newer updationTime", async () => {
|
|
const store = await MLDataStore.open(dir);
|
|
await store.storeFetched(
|
|
new Map([[100, payload([0.5, 0.25, 0.75])]]),
|
|
new Map([[100, 10]]),
|
|
);
|
|
|
|
// 100 is cached and current; 200 has never been fetched.
|
|
expect(
|
|
store.neededFor([
|
|
{ id: 100, updationTime: 10 },
|
|
{ id: 200, updationTime: 5 },
|
|
]),
|
|
).toEqual([200]);
|
|
|
|
// 100's updationTime advanced past what it was fetched at: refetch.
|
|
expect(store.neededFor([{ id: 100, updationTime: 15 }])).toEqual([100]);
|
|
|
|
// Nothing advanced: nothing to fetch.
|
|
expect(store.neededFor([{ id: 100, updationTime: 10 }])).toEqual([]);
|
|
});
|
|
|
|
it("survives a corrupt index without losing the payloads", async () => {
|
|
const store = await MLDataStore.open(dir);
|
|
await store.storeFetched(
|
|
new Map([[100, payload([0.5, 0.25, 0.75])]]),
|
|
new Map([[100, 10]]),
|
|
);
|
|
writeFileSync(join(dir, "clip.json"), "not json");
|
|
|
|
const reopened = await MLDataStore.open(dir);
|
|
expect(reopened.getIndex().fileIDs).toEqual([100]);
|
|
});
|
|
});
|
|
|
|
// --- Library wiring ---------------------------------------------------------
|
|
|
|
const USER_ID = 42;
|
|
const FAST_INTERVAL = 0.02;
|
|
|
|
const collection = (id: number, updationTime: number): Collection => ({
|
|
id,
|
|
ownerID: USER_ID,
|
|
key: new Uint8Array([id & 0xff]),
|
|
name: `album-${id}`,
|
|
type: "album",
|
|
updationTime,
|
|
isShared: false,
|
|
});
|
|
|
|
const file = (
|
|
id: number,
|
|
collectionID: number,
|
|
updationTime: number,
|
|
): EnteFile => ({
|
|
id,
|
|
collectionID,
|
|
ownerID: USER_ID,
|
|
key: new Uint8Array([id & 0xff]),
|
|
metadata: {
|
|
title: `file-${id}.jpg`,
|
|
fileType: "image",
|
|
creationTime: updationTime,
|
|
modificationTime: updationTime,
|
|
},
|
|
file: { decryptionHeader: "aGVhZGVy" },
|
|
thumbnail: { decryptionHeader: "dGh1bWI=" },
|
|
updationTime,
|
|
});
|
|
|
|
// A mock client that serves scripted collection/file pages and per-file ML
|
|
// payloads, recording every ML fetch request so incremental behaviour is
|
|
// provable.
|
|
class MLMockClient {
|
|
userID = USER_ID;
|
|
collectionsQueue: CollectionsPage[] = [];
|
|
filesByCollection = new Map<number, FilesPage[]>();
|
|
mlByFile = new Map<number, MLData>();
|
|
mlFetchCalls: number[][] = [];
|
|
|
|
whoami(): { email: string; userID: number } {
|
|
return { email: "user@example.com", userID: this.userID };
|
|
}
|
|
|
|
async collectionsSince(args: {
|
|
sinceTime: number;
|
|
}): Promise<CollectionsPage> {
|
|
return (
|
|
this.collectionsQueue.shift() ?? {
|
|
collections: [],
|
|
deleted: [],
|
|
cursor: args.sinceTime,
|
|
}
|
|
);
|
|
}
|
|
|
|
async filesSince(args: {
|
|
collectionID: number;
|
|
collectionKey: Uint8Array;
|
|
sinceTime: number;
|
|
}): Promise<FilesPage> {
|
|
const queue = this.filesByCollection.get(args.collectionID);
|
|
return (
|
|
queue?.shift() ?? {
|
|
files: [],
|
|
deleted: [],
|
|
cursor: args.sinceTime,
|
|
}
|
|
);
|
|
}
|
|
|
|
async fetchMLData(args: {
|
|
fileIDs: number[];
|
|
fileKeys: Map<number, Uint8Array>;
|
|
}): Promise<Map<number, MLData>> {
|
|
this.mlFetchCalls.push([...args.fileIDs]);
|
|
const result = new Map<number, MLData>();
|
|
for (const id of args.fileIDs) {
|
|
const p = this.mlByFile.get(id);
|
|
if (p) result.set(id, p);
|
|
}
|
|
return result;
|
|
}
|
|
|
|
filesFor(collectionID: number, ...pages: FilesPage[]): void {
|
|
this.filesByCollection.set(collectionID, pages);
|
|
}
|
|
}
|
|
|
|
describe("Library ML-data fetch on refresh", () => {
|
|
let dir: string;
|
|
let cacheDirectory: string;
|
|
|
|
beforeEach(() => {
|
|
dir = mkdtempSync(join(tmpdir(), "quak-lib-mldata-"));
|
|
cacheDirectory = join(dir, "cache");
|
|
});
|
|
|
|
afterEach(() => {
|
|
rmSync(dir, { recursive: true, force: true });
|
|
});
|
|
|
|
it("fetches, stores and indexes ML data for known files, then is incremental", async () => {
|
|
const client = new MLMockClient();
|
|
client.collectionsQueue.push({
|
|
collections: [collection(1, 100)],
|
|
deleted: [],
|
|
cursor: 100,
|
|
});
|
|
client.filesFor(1, {
|
|
files: [file(1001, 1, 90), file(1002, 1, 95)],
|
|
deleted: [],
|
|
cursor: 95,
|
|
});
|
|
client.mlByFile.set(1001, payload([0.5, 0.25, 0.75]));
|
|
client.mlByFile.set(1002, payload([1, -2, 0.5]));
|
|
|
|
const lib = await Library.open({
|
|
client,
|
|
cacheDirectory,
|
|
refreshIntervalSeconds: FAST_INTERVAL,
|
|
});
|
|
try {
|
|
await vi.waitFor(
|
|
() => {
|
|
expect(lib.status().mlIndexed).toBe(2);
|
|
expect(lib.status().mlStored).toBe(2);
|
|
},
|
|
{ timeout: 2000, interval: 5 },
|
|
);
|
|
|
|
// Both files were fetched, in one batch.
|
|
expect(client.mlFetchCalls.flat().sort((a, b) => a - b)).toEqual([
|
|
1001, 1002,
|
|
]);
|
|
const callsAfterFirst = client.mlFetchCalls.length;
|
|
|
|
// The index is on disk and reloads to the same shape.
|
|
const reopened = await MLDataStore.open(
|
|
join(cacheDirectory, "mldata"),
|
|
);
|
|
expect(reopened.getIndex().fileIDs).toEqual([1001, 1002]);
|
|
|
|
// Later refreshes with nothing new must not refetch.
|
|
await new Promise((r) => setTimeout(r, FAST_INTERVAL * 1000 * 5));
|
|
expect(client.mlFetchCalls.length).toBe(callsAfterFirst);
|
|
} finally {
|
|
lib.close();
|
|
}
|
|
});
|
|
|
|
it("refetches a file whose updationTime advanced", async () => {
|
|
const client = new MLMockClient();
|
|
client.collectionsQueue.push({
|
|
collections: [collection(1, 100)],
|
|
deleted: [],
|
|
cursor: 100,
|
|
});
|
|
client.filesFor(1, {
|
|
files: [file(1001, 1, 90)],
|
|
deleted: [],
|
|
cursor: 90,
|
|
});
|
|
client.mlByFile.set(1001, payload([0.5, 0.25, 0.75]));
|
|
|
|
const lib = await Library.open({
|
|
client,
|
|
cacheDirectory,
|
|
refreshIntervalSeconds: FAST_INTERVAL,
|
|
});
|
|
try {
|
|
await vi.waitFor(() => expect(lib.status().mlIndexed).toBe(1), {
|
|
timeout: 2000,
|
|
interval: 5,
|
|
});
|
|
const callsBefore = client.mlFetchCalls.length;
|
|
|
|
// The file changes on the server (updationTime advances) with a new
|
|
// embedding; the next refresh must refetch it.
|
|
client.mlByFile.set(1001, payload([9, 9, 9]));
|
|
client.collectionsQueue.push({
|
|
collections: [collection(1, 200)],
|
|
deleted: [],
|
|
cursor: 200,
|
|
});
|
|
client.filesFor(1, {
|
|
files: [file(1001, 1, 190)],
|
|
deleted: [],
|
|
cursor: 190,
|
|
});
|
|
|
|
await vi.waitFor(
|
|
() => {
|
|
expect(client.mlFetchCalls.length).toBeGreaterThan(
|
|
callsBefore,
|
|
);
|
|
expect(client.mlFetchCalls.flat()).toContain(1001);
|
|
},
|
|
{ timeout: 2000, interval: 5 },
|
|
);
|
|
|
|
const reopened = await MLDataStore.open(
|
|
join(cacheDirectory, "mldata"),
|
|
);
|
|
expect([...reopened.getIndex().embeddings]).toEqual([9, 9, 9]);
|
|
} finally {
|
|
lib.close();
|
|
}
|
|
});
|
|
});
|