From 5101df2cf2b4022c1511ecc19d263abed931c0d1 Mon Sep 17 00:00:00 2001 From: Rene Zander Date: Fri, 3 Jul 2026 06:07:27 +0000 Subject: [PATCH] feat(hub): resume interrupted downloads via HTTP Range (Node) FileCache now streams large downloads to a deterministic .incomplete file plus an etag/size sidecar, and keeps the partial on failure instead of deleting it. On the next attempt the Node streaming path (return_path) asks getResumeInfo for the offset and sends Range + If-Range: the server appends via 206, or restarts cleanly via 200 if the file changed upstream. A per-key lock preserves the previous concurrent-writer guarantee by falling back to the legacy unique-temp path when the partial is held. Adds FileCache unit tests for resume, truncation, and restart. --- packages/transformers/src/utils/cache.js | 3 + .../transformers/src/utils/cache/FileCache.js | 197 ++++++++++++++++-- packages/transformers/src/utils/hub.js | 40 +++- .../tests/utils/file_cache.test.js | 122 +++++++++++ 4 files changed, 337 insertions(+), 25 deletions(-) create mode 100644 packages/transformers/tests/utils/file_cache.test.js diff --git a/packages/transformers/src/utils/cache.js b/packages/transformers/src/utils/cache.js index ca8e27971..5dc5d84b3 100644 --- a/packages/transformers/src/utils/cache.js +++ b/packages/transformers/src/utils/cache.js @@ -11,6 +11,9 @@ import { CrossOriginStorage } from './cache/CrossOriginStorageCache.js'; * Adds a response to the cache. * @property {(request: string) => Promise} [delete] * Deletes a request from the cache. Returns true if deleted, false otherwise. + * @property {(request: string) => Promise<{size: number, etag: string|null, total: number}|undefined>} [getResumeInfo] + * Optional. Reports a resumable partial download for a request (size on disk, its etag, and expected total), + * enabling the caller to continue an interrupted download via an HTTP `Range` request. */ /** diff --git a/packages/transformers/src/utils/cache/FileCache.js b/packages/transformers/src/utils/cache/FileCache.js index 9e81d4a8e..16a8dfded 100644 --- a/packages/transformers/src/utils/cache/FileCache.js +++ b/packages/transformers/src/utils/cache/FileCache.js @@ -8,9 +8,19 @@ import { apis } from '../../env.js'; // Create a dedicated random instance for generating unique temporary file names const rng = new Random(); +// How long a `.incomplete.lock` may sit before another writer treats it as +// stale (a previous run crashed without releasing it). Well above any realistic +// single-chunk write, but low enough that a resume is not blocked forever. +const LOCK_STALE_MS = 10 * 60 * 1000; + /** * File system cache implementation that implements the CacheInterface. * Provides `match` and `put` methods compatible with the Web Cache API. + * + * Large downloads are streamed to a deterministic `.incomplete` file with + * a small sidecar recording the server `ETag` and expected size. If a download + * is interrupted the partial is kept, so a later attempt can send a `Range` + * request (see `getResumeInfo`) and continue instead of starting from byte 0. */ export class FileCache { /** @@ -37,6 +47,36 @@ export class FileCache { } } + /** + * Report whether a resumable partial download exists for `request`. + * Returns `{ size, etag, total }` when a `.incomplete` file and its sidecar + * are present and internally consistent, otherwise `undefined`. Callers use + * this to send a `Range`/`If-Range` request and continue the download. + * @param {string} request + * @returns {Promise<{size: number, etag: string|null, total: number} | undefined>} + */ + async getResumeInfo(request) { + const incompletePath = path.join(this.path, request) + '.incomplete'; + const sidecarPath = incompletePath + '.json'; + try { + const [stat, sidecarRaw] = await Promise.all([ + fs.promises.stat(incompletePath), + fs.promises.readFile(sidecarPath, 'utf-8'), + ]); + const sidecar = JSON.parse(sidecarRaw); + const total = typeof sidecar.total === 'number' ? sidecar.total : 0; + // Only resume when there is something to resume and we have not + // somehow written past the expected size. Anything inconsistent + // falls through to a clean restart. + if (stat.size > 0 && (total === 0 || stat.size < total)) { + return { size: stat.size, etag: sidecar.etag ?? null, total }; + } + } catch { + // No partial, no sidecar, or unreadable — nothing to resume. + } + return undefined; + } + /** * Adds the given response to the cache. * @param {string} request @@ -47,20 +87,104 @@ export class FileCache { */ async put(request, response, progress_callback = undefined) { const filePath = path.join(this.path, request); + await fs.promises.mkdir(path.dirname(filePath), { recursive: true }); + + const incompletePath = filePath + '.incomplete'; + const sidecarPath = incompletePath + '.json'; + const lockPath = incompletePath + '.lock'; + + // Claim the deterministic partial for this key. If another writer (this + // or another process) already holds it, fall back to the legacy + // random-suffix temp path: correct and non-resumable, but it never + // corrupts the shared partial — preserving the previous concurrency + // guarantee for parallel loads of the same file. + const locked = await this.#acquireLock(lockPath); + if (!locked) { + return this.#putUniqueTemp(filePath, response, progress_callback); + } + + // A 206 (Partial Content) means the server honored our Range request and + // we append to the existing partial. Anything else (200) is a fresh + // download: discard any stale partial and start over. + const resuming = response.status === 206; + const total = this.#expectedTotal(response, resuming); + + let loaded = 0; + if (resuming) { + try { + loaded = (await fs.promises.stat(incompletePath)).size; + } catch { + loaded = 0; + } + } + + try { + // Record the etag + expected total up front so an interrupted write + // still leaves a sidecar the next attempt can validate against. + await fs.promises.writeFile(sidecarPath, JSON.stringify({ etag: response.headers.get('etag'), total })); + + const fileStream = fs.createWriteStream(incompletePath, { + flags: resuming ? 'a' : 'w', + }); + const reader = response.body.getReader(); + + while (true) { + const { done, value } = await reader.read(); + if (done) { + break; + } - // Include both PID and a random suffix so that concurrent put() call within the same process - // (e.g., multiple pipelines loading the same file in parallel) each get their own temp file - // and don't corrupt each other's writes. + await new Promise((resolve, reject) => { + fileStream.write(value, (err) => (err ? reject(err) : resolve())); + }); + + loaded += value.length; + const progress = total ? (loaded / total) * 100 : 0; + + progress_callback?.({ progress, loaded, total }); + } + + await new Promise((resolve, reject) => { + fileStream.close((err) => (err ? reject(err) : resolve())); + }); + + // Guard against a truncated body: if we know the expected size and + // fell short, keep the partial so the next call can resume rather + // than promoting an incomplete file to the final path. + if (total && loaded < total) { + throw new Error(`Incomplete download for "${request}": got ${loaded} of ${total} bytes.`); + } + + // Atomically publish the completed file and drop the sidecar. + await fs.promises.rename(incompletePath, filePath); + await fs.promises.unlink(sidecarPath).catch(() => {}); + } catch (error) { + // Intentionally keep `.incomplete` + sidecar on failure so a later + // attempt can resume from here. Only the lock is released (finally). + throw error; + } finally { + await fs.promises.unlink(lockPath).catch(() => {}); + } + } + + /** + * Legacy write path: stream to a unique temp file and atomically rename. + * Used when the deterministic partial is already locked by a concurrent + * writer, so this call must not touch the shared `.incomplete` file. + * @param {string} filePath + * @param {Response} response + * @param {((data: {progress: number, loaded: number, total: number}) => void) | undefined} progress_callback + * @returns {Promise} + */ + async #putUniqueTemp(filePath, response, progress_callback) { const id = apis.IS_PROCESS_AVAILABLE ? process.pid : Date.now(); const randomSuffix = rng._int32().toString(36); const tmpPath = filePath + `.tmp.${id}.${randomSuffix}`; try { - const contentLength = response.headers.get('Content-Length'); - const total = parseInt(contentLength ?? '0'); + const total = parseInt(response.headers.get('Content-Length') ?? '0'); let loaded = 0; - await fs.promises.mkdir(path.dirname(filePath), { recursive: true }); const fileStream = fs.createWriteStream(tmpPath); const reader = response.body.getReader(); @@ -71,13 +195,7 @@ export class FileCache { } await new Promise((resolve, reject) => { - fileStream.write(value, (err) => { - if (err) { - reject(err); - return; - } - resolve(); - }); + fileStream.write(value, (err) => (err ? reject(err) : resolve())); }); loaded += value.length; @@ -90,12 +208,8 @@ export class FileCache { fileStream.close((err) => (err ? reject(err) : resolve())); }); - // Atomically move the completed temp file to the final path so that - // concurrent readers (other processes or other in-process calls) - // never observe a partially-written file. await fs.promises.rename(tmpPath, filePath); } catch (error) { - // Clean up the temp file if an error occurred during download try { await fs.promises.unlink(tmpPath); } catch {} @@ -103,6 +217,55 @@ export class FileCache { } } + /** + * Total expected file size. For a 206 response it is read from the + * `Content-Range: bytes start-end/total` header; for a fresh 200 from + * `Content-Length`. Returns 0 when the size is unknown. + * @param {Response} response + * @param {boolean} resuming + * @returns {number} + */ + #expectedTotal(response, resuming) { + if (resuming) { + const contentRange = response.headers.get('Content-Range'); + const match = contentRange && /\/(\d+)\s*$/.exec(contentRange); + if (match) { + return parseInt(match[1], 10); + } + } + return parseInt(response.headers.get('Content-Length') ?? '0', 10) || 0; + } + + /** + * Acquire the per-key write lock via an exclusive create (`wx`). If a lock + * already exists but is older than `LOCK_STALE_MS`, it is treated as + * abandoned by a crashed writer and stolen. + * @param {string} lockPath + * @returns {Promise} + */ + async #acquireLock(lockPath) { + try { + const fd = await fs.promises.open(lockPath, 'wx'); + await fd.close(); + return true; + } catch (e) { + if (e.code !== 'EEXIST') { + return false; + } + // Steal a stale lock left behind by a crashed process. + try { + const stat = await fs.promises.stat(lockPath); + if (Date.now() - stat.mtimeMs > LOCK_STALE_MS) { + await fs.promises.unlink(lockPath).catch(() => {}); + const fd = await fs.promises.open(lockPath, 'wx'); + await fd.close(); + return true; + } + } catch {} + return false; + } + } + /** * Deletes the cache entry for the given request. * @param {string} request diff --git a/packages/transformers/src/utils/hub.js b/packages/transformers/src/utils/hub.js index 09f42e168..f0c3292cd 100755 --- a/packages/transformers/src/utils/hub.js +++ b/packages/transformers/src/utils/hub.js @@ -62,9 +62,11 @@ export { MAX_EXTERNAL_DATA_CHUNKS } from './hub/constants.js'; * Helper function to get a file, using either the Fetch API or FileSystem API. * * @param {URL|string} urlOrPath The URL/path of the file to get. + * @param {Record} [extraHeaders] Additional request headers to merge + * on top of the default fetch headers (e.g. `Range`/`If-Range` for resuming a download). * @returns {Promise} A promise that resolves to a FileResponse object (if the file is retrieved using the FileSystem API), or a Response object (if the file is retrieved using the Fetch API). */ -export async function getFile(urlOrPath) { +export async function getFile(urlOrPath, extraHeaders = undefined) { if (env.useFS && !isValidUrl(urlOrPath, ['http:', 'https:', 'blob:'])) { return new FileResponse( urlOrPath instanceof URL @@ -74,9 +76,13 @@ export async function getFile(urlOrPath) { : urlOrPath, ); } else { - return env.fetch(urlOrPath, { - headers: getFetchHeaders(urlOrPath), - }); + const headers = getFetchHeaders(urlOrPath); + if (extraHeaders) { + for (const [key, value] of Object.entries(extraHeaders)) { + headers.set(key, value); + } + } + return env.fetch(urlOrPath, { headers }); } } @@ -330,10 +336,28 @@ export async function loadResourceFile( ); } - // File not found locally, so we try to download it from the remote server - response = await getFile(remoteURL); + // File not found locally, so we try to download it from the remote server. + // + // Resume an interrupted download when the Node file cache still holds a + // partial for this key. This only applies to the streaming path + // (`IS_NODE_ENV && return_path`), where the body is written straight to + // disk; on the buffered path a 206 would yield only the trailing bytes. + // `Range` asks the CDN to continue from the last byte we have; `If-Range` + // makes the server fall back to a full 200 if the file changed upstream. + let resumeHeaders; + if (apis.IS_NODE_ENV && return_path && cache && typeof cache.getResumeInfo === 'function') { + const resume = await cache.getResumeInfo(proposedCacheKey); + if (resume && resume.size > 0) { + resumeHeaders = { Range: `bytes=${resume.size}-` }; + if (resume.etag) { + resumeHeaders['If-Range'] = resume.etag; + } + } + } + response = await getFile(remoteURL, resumeHeaders); - if (response.status !== 200) { + // 200 (full download) and 206 (resumed partial) are both success. + if (response.status !== 200 && response.status !== 206) { return handleError(response.status, remoteURL, fatal); } @@ -346,7 +370,7 @@ export async function loadResourceFile( cache && // 1. A caching system is available typeof Response !== 'undefined' && // 2. `Response` is defined (i.e., we are in a browser-like environment) response instanceof Response && // 3. result is a `Response` object (i.e., not a `FileResponse`) - response.status === 200; // 4. request was successful (status code 200) + (response.status === 200 || response.status === 206); // 4. request was successful (200 full, or 206 resumed) } // Start downloading diff --git a/packages/transformers/tests/utils/file_cache.test.js b/packages/transformers/tests/utils/file_cache.test.js new file mode 100644 index 000000000..67433d1dd --- /dev/null +++ b/packages/transformers/tests/utils/file_cache.test.js @@ -0,0 +1,122 @@ +import fs from "node:fs"; +import os from "node:os"; +import path from "node:path"; + +import { FileCache } from "../../src/utils/cache/FileCache.js"; + +/** + * Build a streaming `Response` from an array of `Uint8Array` chunks. + * @param {Uint8Array[]} chunks + * @param {{status?: number, headers?: Record}} [init] + */ +function streamingResponse(chunks, { status = 200, headers = {} } = {}) { + let i = 0; + const body = new ReadableStream({ + pull(controller) { + if (i >= chunks.length) { + controller.close(); + return; + } + controller.enqueue(chunks[i++]); + }, + }); + return new Response(body, { status, headers }); +} + +const bytes = (str) => new TextEncoder().encode(str); + +describe("FileCache resumable downloads", () => { + let dir; + let cache; + + beforeEach(() => { + dir = fs.mkdtempSync(path.join(os.tmpdir(), "tjs-filecache-")); + cache = new FileCache(dir); + }); + + afterEach(() => { + fs.rmSync(dir, { recursive: true, force: true }); + }); + + it("getResumeInfo returns undefined when there is no partial", async () => { + expect(await cache.getResumeInfo("model.onnx")).toBeUndefined(); + }); + + it("writes a full 200 response and leaves no partial artifacts", async () => { + const content = "hello world"; + await cache.put( + "model.onnx", + streamingResponse([bytes(content)], { + status: 200, + headers: { "Content-Length": String(content.length), etag: '"abc"' }, + }), + ); + + expect(fs.readFileSync(path.join(dir, "model.onnx"), "utf-8")).toBe(content); + expect(fs.existsSync(path.join(dir, "model.onnx.incomplete"))).toBe(false); + expect(fs.existsSync(path.join(dir, "model.onnx.incomplete.json"))).toBe(false); + expect(fs.existsSync(path.join(dir, "model.onnx.incomplete.lock"))).toBe(false); + }); + + it("keeps the partial and sidecar when the body is truncated", async () => { + const key = "big.onnx"; + // Server promises 10 bytes but only delivers 4. + await expect( + cache.put( + key, + streamingResponse([bytes("0123")], { + status: 200, + headers: { "Content-Length": "10", etag: '"v1"' }, + }), + ), + ).rejects.toThrow(/Incomplete download/); + + // Final file must NOT be published; partial + sidecar remain for resume. + expect(fs.existsSync(path.join(dir, key))).toBe(false); + expect(fs.readFileSync(path.join(dir, key + ".incomplete"), "utf-8")).toBe("0123"); + expect(fs.existsSync(path.join(dir, key + ".incomplete.lock"))).toBe(false); + + const resume = await cache.getResumeInfo(key); + expect(resume).toEqual({ size: 4, etag: '"v1"', total: 10 }); + }); + + it("appends a 206 partial-content response onto the existing partial", async () => { + const key = "weights.onnx"; + const full = "0123456789"; + + // Seed a prior interrupted download: first 4 bytes + sidecar. + fs.writeFileSync(path.join(dir, key + ".incomplete"), full.slice(0, 4)); + fs.writeFileSync(path.join(dir, key + ".incomplete.json"), JSON.stringify({ etag: '"v1"', total: 10 })); + + // Server continues from byte 4 with a 206. + await cache.put( + key, + streamingResponse([bytes(full.slice(4))], { + status: 206, + headers: { "Content-Range": "bytes 4-9/10", "Content-Length": "6", etag: '"v1"' }, + }), + ); + + expect(fs.readFileSync(path.join(dir, key), "utf-8")).toBe(full); + expect(fs.existsSync(path.join(dir, key + ".incomplete"))).toBe(false); + expect(fs.existsSync(path.join(dir, key + ".incomplete.json"))).toBe(false); + expect(await cache.getResumeInfo(key)).toBeUndefined(); + }); + + it("restarts from scratch on a 200 even if a stale partial exists", async () => { + const key = "restart.onnx"; + fs.writeFileSync(path.join(dir, key + ".incomplete"), "STALE-JUNK"); + fs.writeFileSync(path.join(dir, key + ".incomplete.json"), JSON.stringify({ etag: '"old"', total: 10 })); + + const content = "freshdata!"; + await cache.put( + key, + streamingResponse([bytes(content)], { + status: 200, + headers: { "Content-Length": String(content.length), etag: '"new"' }, + }), + ); + + expect(fs.readFileSync(path.join(dir, key), "utf-8")).toBe(content); + }); +});