clawbot
d7f415fe29
check / check (push) Successful in 14s
Adds the machine-learning (magic) data layer: fetches per-file ML payloads (face detections + CLIP embeddings) via the existing metadata-backup fetch through the metadata pool after each refresh, decrypts and gunzips them, and stores one mldata/<fileID>.json per file by rename (present-means-complete). A derived index (mldata/clip.f32 + clip.json) loads in one read and is rebuilt whenever it disagrees with the payloads on disk in either direction, so an interrupted backfill self-heals. Never in metadata.json; incremental on later refreshes; progress via onProgress/status. Model: opus-4-8
379 lines
14 KiB
TypeScript
379 lines
14 KiB
TypeScript
// The on-disk cache of Ente's per-file machine-learning data and the CLIP
|
|
// index derived from it (issue #49).
|
|
//
|
|
// Under `<cacheDirectory>/mldata/` this keeps:
|
|
//
|
|
// - `<fileID>.json` — one decrypted, gunzipped payload per file, written by
|
|
// rename. Its presence means it is complete: a torn write never leaves a
|
|
// half-file, so the set of these files is the source of truth for what is
|
|
// cached. The full payload (face boxes, landmarks, embeddings) is read back
|
|
// from here on demand and never held in RAM.
|
|
//
|
|
// - `clip.f32` + `clip.json` — the derived index the content search runs on.
|
|
// `clip.json` lists the indexed fileIDs in order plus the embedding length;
|
|
// `clip.f32` is those CLIP embeddings packed as one `Float32Array`, so the
|
|
// index loads in a single read with no per-vector parse. The index is
|
|
// rebuilt from the payloads whenever it is missing or structurally
|
|
// disagrees with the files present, and appended to as new payloads arrive.
|
|
//
|
|
// - `fetched.json` — a small map of fileID to the `updationTime` it was
|
|
// fetched at. This is best-effort bookkeeping for refetch decisions (a file
|
|
// whose `updationTime` later advances is refetched); the payloads, not this
|
|
// file, remain the record of what is cached, so losing it only forgoes
|
|
// update-driven refetch until the next fetch rewrites it.
|
|
//
|
|
// In RAM this holds only the id list and the packed `Float32Array`.
|
|
|
|
import { mkdir, readFile, readdir } from "node:fs/promises";
|
|
import { join } from "node:path";
|
|
|
|
import { writeAtomic } from "../download/index.js";
|
|
import type { MLData } from "../mldata-fetch.js";
|
|
|
|
const CLIP_VECTORS = "clip.f32";
|
|
const CLIP_INDEX = "clip.json";
|
|
const FETCHED = "fetched.json";
|
|
// A payload file is named for its fileID alone; the derived files above are
|
|
// not, so this pattern picks out payloads and nothing else.
|
|
const PAYLOAD_RE = /^(\d+)\.json$/;
|
|
const BYTES_PER_FLOAT = 4;
|
|
|
|
// The on-disk form of `clip.json`.
|
|
interface ClipIndexFile {
|
|
fileIDs: number[];
|
|
embeddingLength: number;
|
|
}
|
|
|
|
// A file the model knows about, for deciding what to fetch.
|
|
export interface MLDataFile {
|
|
id: number;
|
|
updationTime: number;
|
|
}
|
|
|
|
// The RAM index the search reads: `fileIDs[i]` owns the `embeddingLength`
|
|
// floats of `embeddings` starting at `i * embeddingLength`.
|
|
export interface MLIndex {
|
|
fileIDs: number[];
|
|
embeddingLength: number;
|
|
embeddings: Float32Array;
|
|
}
|
|
|
|
// Pull the CLIP embedding out of a payload, or undefined when it is absent or
|
|
// misshapen. Kept strict so a bad payload is skipped rather than corrupting the
|
|
// packed index.
|
|
const clipEmbedding = (payload: MLData): number[] | undefined => {
|
|
const clip = payload.clip;
|
|
if (typeof clip !== "object" || clip === null) return undefined;
|
|
const embedding = (clip as { embedding?: unknown }).embedding;
|
|
if (!Array.isArray(embedding)) return undefined;
|
|
if (embedding.some((v) => typeof v !== "number" || !Number.isFinite(v)))
|
|
return undefined;
|
|
return embedding as number[];
|
|
};
|
|
|
|
export class MLDataStore {
|
|
readonly dir: string;
|
|
|
|
// fileIDs whose payload JSON is present on disk (present means complete).
|
|
private readonly present = new Set<number>();
|
|
// fileID -> updationTime it was fetched at.
|
|
private readonly fetched = new Map<number, number>();
|
|
|
|
// The packed index and where each id sits in it.
|
|
private ids: number[] = [];
|
|
private embeddingLength = 0;
|
|
private embeddings = new Float32Array(0);
|
|
private readonly pos = new Map<number, number>();
|
|
|
|
private constructor(dir: string) {
|
|
this.dir = dir;
|
|
}
|
|
|
|
// Open (creating the directory) and load the id list and packed index into
|
|
// RAM, rebuilding the index from the payloads when it is missing or does
|
|
// not match the files present.
|
|
static async open(dir: string): Promise<MLDataStore> {
|
|
const store = new MLDataStore(dir);
|
|
await mkdir(dir, { recursive: true });
|
|
await store.loadPresent();
|
|
await store.loadFetched();
|
|
if (!(await store.tryLoadIndex())) await store.rebuildIndex();
|
|
return store;
|
|
}
|
|
|
|
// The fileIDs among `files` that must be fetched: every file with no
|
|
// payload yet (first run, then new files), plus any whose `updationTime`
|
|
// has advanced past the one its cached payload was fetched at. Returned
|
|
// sorted and unique.
|
|
neededFor(files: MLDataFile[]): number[] {
|
|
const latest = new Map<number, number>();
|
|
for (const f of files) {
|
|
const seen = latest.get(f.id);
|
|
if (seen === undefined || f.updationTime > seen)
|
|
latest.set(f.id, f.updationTime);
|
|
}
|
|
const needed: number[] = [];
|
|
for (const [id, updationTime] of latest) {
|
|
if (!this.present.has(id)) {
|
|
needed.push(id);
|
|
continue;
|
|
}
|
|
const at = this.fetched.get(id);
|
|
if (at !== undefined && updationTime > at) needed.push(id);
|
|
}
|
|
return needed.sort((a, b) => a - b);
|
|
}
|
|
|
|
// Store a batch of fetched payloads: write one file per id, fold their CLIP
|
|
// embeddings into the packed index (in place for a refetch, appended for a
|
|
// new file), and persist the derived files. Returns how many payloads were
|
|
// stored and how many ids the index now holds.
|
|
async storeFetched(
|
|
payloads: Map<number, MLData>,
|
|
updation: Map<number, number>,
|
|
): Promise<{ stored: number; indexed: number }> {
|
|
if (payloads.size === 0) return { stored: 0, indexed: this.ids.length };
|
|
|
|
for (const [id, payload] of payloads) {
|
|
await this.writePayload(id, payload);
|
|
this.present.add(id);
|
|
const at = updation.get(id);
|
|
if (at !== undefined) this.fetched.set(id, at);
|
|
}
|
|
|
|
const updates: { at: number; vector: number[] }[] = [];
|
|
const appends: { id: number; vector: number[] }[] = [];
|
|
for (const [id, payload] of payloads) {
|
|
const vector = clipEmbedding(payload);
|
|
if (!vector) continue;
|
|
if (this.embeddingLength === 0 && this.ids.length === 0)
|
|
this.embeddingLength = vector.length;
|
|
// The index is fixed-width; a vector of another length (never seen
|
|
// from Ente's CLIP model) is stored but left out of the index.
|
|
if (vector.length !== this.embeddingLength) continue;
|
|
const at = this.pos.get(id);
|
|
if (at !== undefined) updates.push({ at, vector });
|
|
else appends.push({ id, vector });
|
|
}
|
|
|
|
for (const { at, vector } of updates)
|
|
this.embeddings.set(vector, at * this.embeddingLength);
|
|
|
|
if (appends.length > 0) {
|
|
const length = this.embeddingLength;
|
|
const grown = new Float32Array(
|
|
this.embeddings.length + appends.length * length,
|
|
);
|
|
grown.set(this.embeddings);
|
|
let offset = this.embeddings.length;
|
|
for (const { id, vector } of appends) {
|
|
grown.set(vector, offset);
|
|
this.pos.set(id, this.ids.length);
|
|
this.ids.push(id);
|
|
offset += length;
|
|
}
|
|
this.embeddings = grown;
|
|
}
|
|
|
|
await this.persistIndex();
|
|
await this.persistFetched();
|
|
return { stored: payloads.size, indexed: this.ids.length };
|
|
}
|
|
|
|
// The packed index the search runs on. The id list is copied so callers
|
|
// cannot disturb the store's own order; the embeddings are the live buffer.
|
|
getIndex(): MLIndex {
|
|
return {
|
|
fileIDs: [...this.ids],
|
|
embeddingLength: this.embeddingLength,
|
|
embeddings: this.embeddings,
|
|
};
|
|
}
|
|
|
|
// The full payload for a file, read from disk, or undefined when it is not
|
|
// cached or does not parse.
|
|
async readPayload(fileID: number): Promise<MLData | undefined> {
|
|
if (!this.present.has(fileID)) return undefined;
|
|
let raw: string;
|
|
try {
|
|
raw = await readFile(this.payloadPath(fileID), "utf-8");
|
|
} catch {
|
|
return undefined;
|
|
}
|
|
try {
|
|
return JSON.parse(raw) as MLData;
|
|
} catch {
|
|
return undefined;
|
|
}
|
|
}
|
|
|
|
stats(): { stored: number; indexed: number } {
|
|
return { stored: this.present.size, indexed: this.ids.length };
|
|
}
|
|
|
|
private payloadPath(id: number): string {
|
|
return join(this.dir, `${id}.json`);
|
|
}
|
|
|
|
private async writePayload(id: number, payload: MLData): Promise<void> {
|
|
await writeAtomic(
|
|
this.payloadPath(id),
|
|
new TextEncoder().encode(JSON.stringify(payload)),
|
|
);
|
|
}
|
|
|
|
private async loadPresent(): Promise<void> {
|
|
let names: string[];
|
|
try {
|
|
names = await readdir(this.dir);
|
|
} catch {
|
|
return;
|
|
}
|
|
for (const name of names) {
|
|
const match = PAYLOAD_RE.exec(name);
|
|
if (match) this.present.add(Number(match[1]));
|
|
}
|
|
}
|
|
|
|
private async loadFetched(): Promise<void> {
|
|
let raw: string;
|
|
try {
|
|
raw = await readFile(join(this.dir, FETCHED), "utf-8");
|
|
} catch {
|
|
return;
|
|
}
|
|
try {
|
|
const parsed = JSON.parse(raw) as Record<string, unknown>;
|
|
for (const [key, value] of Object.entries(parsed)) {
|
|
const id = Number(key);
|
|
if (
|
|
Number.isInteger(id) &&
|
|
typeof value === "number" &&
|
|
this.present.has(id)
|
|
)
|
|
this.fetched.set(id, value);
|
|
}
|
|
} catch {
|
|
// Corrupt bookkeeping degrades refetch decisions, never fails open.
|
|
}
|
|
}
|
|
|
|
// Load the packed index if it is present and agrees with the payloads in
|
|
// both directions: every id it names must still be present, its vector file
|
|
// must be exactly the size the id count and embedding length imply, and no
|
|
// embedding-bearing payload on disk may be missing from it. Returns whether
|
|
// it loaded.
|
|
private async tryLoadIndex(): Promise<boolean> {
|
|
let metaRaw: string;
|
|
try {
|
|
metaRaw = await readFile(join(this.dir, CLIP_INDEX), "utf-8");
|
|
} catch {
|
|
return false;
|
|
}
|
|
let meta: ClipIndexFile;
|
|
try {
|
|
meta = JSON.parse(metaRaw) as ClipIndexFile;
|
|
} catch {
|
|
return false;
|
|
}
|
|
if (
|
|
!Array.isArray(meta.fileIDs) ||
|
|
typeof meta.embeddingLength !== "number"
|
|
)
|
|
return false;
|
|
if (meta.fileIDs.some((id) => !this.present.has(id))) return false;
|
|
|
|
// The reverse must hold too. A payload carrying an embedding but absent
|
|
// from the index means the index is stale — realistically the process
|
|
// died after storeFetched renamed the payloads into place but before it
|
|
// rewrote clip.json/clip.f32. Loading such an index as "consistent"
|
|
// would drop those embeddings for good (neededFor sees the payloads
|
|
// present and never refetches), so treat it as a disagreement and
|
|
// rebuild. Only present ids the index omits are read; a payload
|
|
// legitimately without an embedding stays out and forces no rebuild.
|
|
const indexed = new Set(meta.fileIDs);
|
|
for (const id of this.present) {
|
|
if (indexed.has(id)) continue;
|
|
const payload = await this.readPayload(id);
|
|
if (payload && clipEmbedding(payload)) return false;
|
|
}
|
|
|
|
let bytes: Buffer;
|
|
try {
|
|
bytes = await readFile(join(this.dir, CLIP_VECTORS));
|
|
} catch {
|
|
return false;
|
|
}
|
|
const expected =
|
|
meta.fileIDs.length * meta.embeddingLength * BYTES_PER_FLOAT;
|
|
if (bytes.byteLength !== expected) return false;
|
|
|
|
// One read, no parse: copy into an aligned buffer and view it as
|
|
// floats. The copy is needed because a Buffer from the pool can start
|
|
// at an offset a Float32Array cannot be laid over.
|
|
const aligned = new Uint8Array(bytes.byteLength);
|
|
aligned.set(bytes);
|
|
this.embeddings = new Float32Array(aligned.buffer);
|
|
this.embeddingLength = meta.embeddingLength;
|
|
this.ids = [...meta.fileIDs];
|
|
this.pos.clear();
|
|
this.ids.forEach((id, i) => this.pos.set(id, i));
|
|
return true;
|
|
}
|
|
|
|
// Rebuild the packed index by reading every payload present, then persist
|
|
// it. Payloads without a CLIP embedding (or of an unexpected length) are
|
|
// simply not indexed.
|
|
private async rebuildIndex(): Promise<void> {
|
|
this.ids = [];
|
|
this.pos.clear();
|
|
this.embeddingLength = 0;
|
|
const vectors: number[][] = [];
|
|
for (const id of [...this.present].sort((a, b) => a - b)) {
|
|
const payload = await this.readPayload(id);
|
|
if (!payload) continue;
|
|
const vector = clipEmbedding(payload);
|
|
if (!vector) continue;
|
|
if (this.embeddingLength === 0)
|
|
this.embeddingLength = vector.length;
|
|
if (vector.length !== this.embeddingLength) continue;
|
|
this.pos.set(id, this.ids.length);
|
|
this.ids.push(id);
|
|
vectors.push(vector);
|
|
}
|
|
const length = this.embeddingLength;
|
|
const packed = new Float32Array(this.ids.length * length);
|
|
vectors.forEach((vector, i) => packed.set(vector, i * length));
|
|
this.embeddings = packed;
|
|
await this.persistIndex();
|
|
}
|
|
|
|
private async persistIndex(): Promise<void> {
|
|
const meta: ClipIndexFile = {
|
|
fileIDs: this.ids,
|
|
embeddingLength: this.embeddingLength,
|
|
};
|
|
await writeAtomic(
|
|
join(this.dir, CLIP_INDEX),
|
|
new TextEncoder().encode(JSON.stringify(meta)),
|
|
);
|
|
await writeAtomic(
|
|
join(this.dir, CLIP_VECTORS),
|
|
new Uint8Array(
|
|
this.embeddings.buffer,
|
|
this.embeddings.byteOffset,
|
|
this.embeddings.byteLength,
|
|
),
|
|
);
|
|
}
|
|
|
|
private async persistFetched(): Promise<void> {
|
|
const record: Record<string, number> = {};
|
|
for (const [id, at] of this.fetched) record[id] = at;
|
|
await writeAtomic(
|
|
join(this.dir, FETCHED),
|
|
new TextEncoder().encode(JSON.stringify(record)),
|
|
);
|
|
}
|
|
}
|