Add a Cache-API model file downloader that reports load progress and caches only complete files
Co-Authored-By: Claude Opus 5.5 <[email protected]>
This commit is contained in:
Lucas Winther
co-authored by
Claude Opus 5.5
Claude Opus 5.5
parent
c9fdab1726
commit
827e9a8dae
@@ -0,0 +1,121 @@
|
||||
/**
|
||||
* Weight-file download + cache for engine paths that don't go through
|
||||
* Transformers.js's own loader (Parakeet, D-WS7a).
|
||||
*
|
||||
* Files land in the same Cache Storage cache Transformers.js uses
|
||||
* (`transformers-cache`, D-005), keyed by their Hub URL, so the downloaded
|
||||
* badge, the storage panel, and delete (src/lib/modelCache.ts) work for them
|
||||
* unchanged. Progress is reported as the same `load-progress` messages the
|
||||
* Whisper path forwards, so ModelLoadProgress needs no second code path.
|
||||
*
|
||||
* A file is `put` only once its whole body has arrived: a cancelled or
|
||||
* failed download (worker terminated, network drop) leaves nothing behind
|
||||
* that could later show as "available offline".
|
||||
*
|
||||
* Dependencies are injected so `bun test` runs this without Cache Storage.
|
||||
*/
|
||||
|
||||
import { TRANSFORMERS_CACHE_NAME } from "../lib/modelCache";
|
||||
import type { LoadProgressMessage } from "./protocol";
|
||||
|
||||
/** The slice of `Cache` used here. */
|
||||
export interface ModelFileCache {
|
||||
match(url: string): Promise<Response | undefined>;
|
||||
put(url: string, response: Response): Promise<void>;
|
||||
}
|
||||
|
||||
/** The slice of `CacheStorage` used here. */
|
||||
export interface ModelFileCacheStorage {
|
||||
open(name: string): Promise<ModelFileCache>;
|
||||
}
|
||||
|
||||
export interface FetchModelFileDeps {
|
||||
/** Cache Storage, or undefined where it doesn't exist (download uncached). */
|
||||
caches: ModelFileCacheStorage | undefined;
|
||||
fetch: (url: string) => Promise<Response>;
|
||||
onProgress: (message: LoadProgressMessage) => void;
|
||||
}
|
||||
|
||||
/** Fetch one model file, from cache when present, else download and cache it. */
|
||||
export async function fetchModelFile(
|
||||
url: string,
|
||||
file: string,
|
||||
deps: FetchModelFileDeps,
|
||||
): Promise<Uint8Array> {
|
||||
deps.onProgress({ type: "load-progress", status: "initiate", file });
|
||||
const cache = deps.caches ? await deps.caches.open(TRANSFORMERS_CACHE_NAME) : undefined;
|
||||
|
||||
const cached = await cache?.match(url);
|
||||
if (cached) {
|
||||
const bytes = new Uint8Array(await cached.arrayBuffer());
|
||||
deps.onProgress({ type: "load-progress", status: "done", file });
|
||||
return bytes;
|
||||
}
|
||||
|
||||
deps.onProgress({ type: "load-progress", status: "download", file });
|
||||
const response = await deps.fetch(url);
|
||||
if (!response.ok || !response.body) {
|
||||
throw new Error(`Could not download ${file} (HTTP ${response.status})`);
|
||||
}
|
||||
const total = Number(response.headers.get("Content-Length")) || 0;
|
||||
const bytes = await readWithProgress(response.body, total, (loaded) => {
|
||||
const size = total || loaded;
|
||||
deps.onProgress({
|
||||
type: "load-progress",
|
||||
status: "progress",
|
||||
file,
|
||||
loaded,
|
||||
total: size,
|
||||
progress: (loaded / size) * 100,
|
||||
});
|
||||
});
|
||||
|
||||
if (cache) {
|
||||
try {
|
||||
await cache.put(
|
||||
url,
|
||||
new Response(bytes as Uint8Array<ArrayBuffer>, {
|
||||
headers: { "Content-Length": String(bytes.length) },
|
||||
}),
|
||||
);
|
||||
} catch {
|
||||
// Quota exceeded / private mode: the model still runs this session, it
|
||||
// just won't be available offline.
|
||||
}
|
||||
}
|
||||
deps.onProgress({ type: "load-progress", status: "done", file });
|
||||
return bytes;
|
||||
}
|
||||
|
||||
/**
|
||||
* Read a body into one buffer — preallocated when the length is known, so a
|
||||
* 650 MB encoder isn't also held as a list of chunks. Reports progress at
|
||||
* most once per whole percent (a large file arrives in thousands of reads).
|
||||
*/
|
||||
async function readWithProgress(
|
||||
body: ReadableStream<Uint8Array>,
|
||||
total: number,
|
||||
onLoaded: (loaded: number) => void,
|
||||
): Promise<Uint8Array> {
|
||||
let buffer = new Uint8Array(total);
|
||||
let loaded = 0;
|
||||
let lastPercent = -1;
|
||||
const reader = body.getReader();
|
||||
for (;;) {
|
||||
const { done, value } = await reader.read();
|
||||
if (done) break;
|
||||
if (loaded + value.length > buffer.length) {
|
||||
const grown = new Uint8Array(Math.max(buffer.length * 2, loaded + value.length));
|
||||
grown.set(buffer.subarray(0, loaded));
|
||||
buffer = grown;
|
||||
}
|
||||
buffer.set(value, loaded);
|
||||
loaded += value.length;
|
||||
const percent = total ? Math.floor((loaded / total) * 100) : loaded;
|
||||
if (percent !== lastPercent) {
|
||||
lastPercent = percent;
|
||||
onLoaded(loaded);
|
||||
}
|
||||
}
|
||||
return loaded === buffer.length ? buffer : buffer.slice(0, loaded);
|
||||
}
|
||||
@@ -0,0 +1,121 @@
|
||||
import { describe, expect, test } from "bun:test";
|
||||
import { fetchModelFile, type ModelFileCacheStorage } from "../src/worker/modelFiles";
|
||||
import type { LoadProgressMessage } from "../src/worker/protocol";
|
||||
|
||||
// Cache-API-backed weight download for the Parakeet path (D-WS7a, D-005).
|
||||
|
||||
const URL_A = "https://huggingface.co/org/model/resolve/rev/encoder.onnx";
|
||||
|
||||
function fakeCaches(initial: Record<string, Uint8Array> = {}) {
|
||||
const stored = new Map<string, Uint8Array>(Object.entries(initial));
|
||||
const opened: string[] = [];
|
||||
const storage: ModelFileCacheStorage = {
|
||||
async open(name) {
|
||||
opened.push(name);
|
||||
return {
|
||||
async match(url) {
|
||||
const bytes = stored.get(url);
|
||||
return bytes ? new Response(bytes as Uint8Array<ArrayBuffer>) : undefined;
|
||||
},
|
||||
async put(url, response) {
|
||||
stored.set(url, new Uint8Array(await response.arrayBuffer()));
|
||||
},
|
||||
};
|
||||
},
|
||||
};
|
||||
return { storage, stored, opened };
|
||||
}
|
||||
|
||||
function streamingFetch(bytes: Uint8Array, pieces: number, status = 200) {
|
||||
const calls: string[] = [];
|
||||
const fetchFn = async (url: string) => {
|
||||
calls.push(url);
|
||||
if (status !== 200) return new Response("nope", { status });
|
||||
const size = Math.ceil(bytes.length / pieces);
|
||||
const body = new ReadableStream<Uint8Array>({
|
||||
start(controller) {
|
||||
for (let i = 0; i < bytes.length; i += size) controller.enqueue(bytes.slice(i, i + size));
|
||||
controller.close();
|
||||
},
|
||||
});
|
||||
return new Response(body, { headers: { "Content-Length": String(bytes.length) } });
|
||||
};
|
||||
return { fetchFn, calls };
|
||||
}
|
||||
|
||||
describe("fetchModelFile", () => {
|
||||
test("a cached file is returned without touching the network", async () => {
|
||||
const cached = new Uint8Array([1, 2, 3]);
|
||||
const { storage, opened } = fakeCaches({ [URL_A]: cached });
|
||||
const { fetchFn, calls } = streamingFetch(new Uint8Array(), 1);
|
||||
const events: LoadProgressMessage[] = [];
|
||||
const bytes = await fetchModelFile(URL_A, "encoder.onnx", {
|
||||
caches: storage,
|
||||
fetch: fetchFn,
|
||||
onProgress: (event) => events.push(event),
|
||||
});
|
||||
expect([...bytes]).toEqual([1, 2, 3]);
|
||||
expect(calls).toEqual([]);
|
||||
expect(opened).toEqual(["transformers-cache"]);
|
||||
expect(events.map((event) => event.status)).toEqual(["initiate", "done"]);
|
||||
});
|
||||
|
||||
test("a miss downloads with byte progress and stores the complete file", async () => {
|
||||
const payload = new Uint8Array(1000).map((_, i) => i % 251);
|
||||
const { storage, stored } = fakeCaches();
|
||||
const { fetchFn, calls } = streamingFetch(payload, 4);
|
||||
const events: LoadProgressMessage[] = [];
|
||||
const bytes = await fetchModelFile(URL_A, "encoder.onnx", {
|
||||
caches: storage,
|
||||
fetch: fetchFn,
|
||||
onProgress: (event) => events.push(event),
|
||||
});
|
||||
expect(calls).toEqual([URL_A]);
|
||||
expect([...bytes]).toEqual([...payload]);
|
||||
expect([...stored.get(URL_A)!]).toEqual([...payload]);
|
||||
const progress = events.filter((event) => event.status === "progress");
|
||||
expect(progress.length).toBe(4);
|
||||
const last = progress[progress.length - 1]!;
|
||||
expect(last).toMatchObject({ file: "encoder.onnx", loaded: 1000, total: 1000, progress: 100 });
|
||||
expect(events[0]!.status).toBe("initiate");
|
||||
expect(events[events.length - 1]!.status).toBe("done");
|
||||
});
|
||||
|
||||
test("an HTTP error throws a plain message and stores nothing", async () => {
|
||||
const { storage, stored } = fakeCaches();
|
||||
const { fetchFn } = streamingFetch(new Uint8Array(), 1, 404);
|
||||
await expect(
|
||||
fetchModelFile(URL_A, "encoder.onnx", { caches: storage, fetch: fetchFn, onProgress: () => {} }),
|
||||
).rejects.toThrow("Could not download encoder.onnx (HTTP 404)");
|
||||
expect(stored.size).toBe(0);
|
||||
});
|
||||
|
||||
test("an interrupted download stores nothing (no false 'available offline')", async () => {
|
||||
const { storage, stored } = fakeCaches();
|
||||
const fetchFn = async () =>
|
||||
new Response(
|
||||
new ReadableStream<Uint8Array>({
|
||||
start(controller) {
|
||||
controller.enqueue(new Uint8Array(10));
|
||||
controller.error(new Error("network down"));
|
||||
},
|
||||
}),
|
||||
{ headers: { "Content-Length": "100" } },
|
||||
);
|
||||
await expect(
|
||||
fetchModelFile(URL_A, "encoder.onnx", { caches: storage, fetch: fetchFn, onProgress: () => {} }),
|
||||
).rejects.toThrow();
|
||||
expect(stored.size).toBe(0);
|
||||
});
|
||||
|
||||
test("without Cache Storage it still downloads", async () => {
|
||||
const { fetchFn, calls } = streamingFetch(new Uint8Array([9]), 1);
|
||||
const bytes = await fetchModelFile(URL_A, "encoder.onnx", {
|
||||
caches: undefined,
|
||||
fetch: fetchFn,
|
||||
onProgress: () => {},
|
||||
});
|
||||
expect([...bytes]).toEqual([9]);
|
||||
expect(calls).toEqual([URL_A]);
|
||||
});
|
||||
});
|
||||
Reference in New Issue
Block a user