diff --git a/server/lib/downloadPreflight.js b/server/lib/downloadPreflight.js index 4c60db550a..7ffb395c13 100644 --- a/server/lib/downloadPreflight.js +++ b/server/lib/downloadPreflight.js @@ -542,8 +542,7 @@ const IDLE_STALL_TIMEOUT_MS = 20 * 60 * 1000; * migrated download path stayed unprotected until somebody remembered to add * that clause. Registering here removes that step for every path that claims a * slot; what such a path must still do is BE LOADED, which - * `orphanedPartialGc.test.js` pins. (A path that claims no slot — `loras.js` - * today — is still unprotected; issue #6190 migrates it.) + * `orphanedPartialGc.test.js` pins. * * Call `createDownloadSlot` at MODULE scope only: this Set is add-only, so a * per-request slot would grow it for the process's lifetime and slow diff --git a/server/services/loras.js b/server/services/loras.js index 1e72456f74..0ce404428c 100644 --- a/server/services/loras.js +++ b/server/services/loras.js @@ -23,6 +23,7 @@ import { ServerError } from '../lib/errorHandler.js'; import { assessDownloadPreflight, assertDownloadFits, + createDownloadSlot, etagPathFor, probeRemoteSize, siblingDownloadMeta, @@ -82,6 +83,14 @@ const loraMetadataCache = new Map(); // files, while two explicitly different versions remain independent. const civitaiInstalls = createSingleFlight(); +// Keyed on the resolved destination path — what `isAnyDownloadInFlight` (and the +// orphaned-partial GC sweep) asks about, and what a second install of the same +// file would otherwise race. Not `exclusive`: LoRA files are small relative to a +// checkpoint, so installing several at once is normal, not a hazard. Cancel +// discards the `.partial` (`keepPartialOnCancel` defaults false) — a single +// abandoned file, not a multi-shard checkpoint worth resuming. +const downloadSlot = createDownloadSlot({ codePrefix: 'LORA' }); + const sidecarPath = (loraFilename) => join(PATHS.loras, `${loraFilename}${SIDECAR_SUFFIX}`); const invalidateLoraMetadataCache = (filename) => { loraMetadataCache.delete(join(PATHS.loras, filename)); @@ -465,70 +474,88 @@ const downloadToFile = async (url, destPath, { fetchImpl = fetch, headers = {} , throw new ServerError(`${label} download failed: ${res.status} ${res.statusText}`, { status: 502, code }); }; - let lastEmit = 0; - let lastTick = { received: 0, total: 0 }; - const { tmpPath } = await streamResumableDownload({ - url, - destPath, - headers, - fetchImpl, - signal, - finalize: false, - isCancelled: () => Boolean(signal?.aborted), - onHttpError, - onBytes: onProgress - ? (received, total) => { - lastTick = { received, total }; - const now = Date.now(); - if (now - lastEmit < 150) return; - lastEmit = now; - onProgress({ received, total }); - } - : undefined, + // Claim before the first await: two presses landing inside the same tick + // would otherwise both pass the caller's existsSync pre-check and start a + // parallel transfer of the same file. Released in the finally below. + const slot = downloadSlot.claim(destPath, { + busyMessage: `${basename(destPath)} is already downloading`, }); - // `finalize: false` means streamResumableDownload never had a chance to - // clean up its own etag sidecar (that only happens on ITS finalize path). - // Reaching here means the stream completed successfully — every branch - // below either moves or deletes tmpPath outright, never leaves it for a - // future resume, so the sidecar describing it is equally done. - await rmGuarded(etagPathFor(destPath), { force: true }).catch(() => {}); - if (onProgress) onProgress(lastTick); - // Atomic no-clobber finalize: `link` is POSIX-atomic and fails with EEXIST - // when destPath already exists (concurrent install that snuck past our - // pre-check). On success we unlink the tmp; on EEXIST we clean up and - // throw CIVITAI_ALREADY_INSTALLED. For other link errors (cross-device - // EXDEV, read-only fs, etc.) fall back to rename, which is the only - // portable option on those platforms. - const linkErr = await link(tmpPath, destPath).catch((e) => e); - if (!linkErr) { - await unlinkGuarded(tmpPath).catch(() => {}); - return; + // An external abort signal (SSE client-disconnect) cancels through the slot + // rather than its own controller, so the typed CANCELLED/STALLED error from + // `wrapError` below applies uniformly regardless of who triggered the abort. + const onExternalAbort = () => downloadSlot.cancel(destPath); + if (signal) { + if (signal.aborted) onExternalAbort(); + else signal.addEventListener('abort', onExternalAbort, { once: true }); } - if (linkErr.code === 'EEXIST') { - await rmGuarded(tmpPath, { force: true }).catch(() => {}); - const basename_ = basename(destPath); - throw new ServerError( - `Already installed: ${basename_}. Delete it first or pick a different version.`, - { status: 409, code: 'CIVITAI_ALREADY_INSTALLED' }, - ); - } - // EXDEV or similar — fall back to rename. Re-check destPath right before - // the rename so a concurrent install that landed between our link attempt - // and now can't be silently clobbered (POSIX rename overwrites). Treat - // late-arriving dest as CIVITAI_ALREADY_INSTALLED, matching the EEXIST - // path above. - if (existsSync(destPath)) { - await rmGuarded(tmpPath, { force: true }).catch(() => {}); - const basename_ = basename(destPath); - throw new ServerError( - `Already installed: ${basename_}. Delete it first or pick a different version.`, - { status: 409, code: 'CIVITAI_ALREADY_INSTALLED' }, - ); + const emitProgress = onProgress + ? slot.throttle((received, total) => onProgress({ received, total })) + : null; + try { + let lastTick = { received: 0, total: 0 }; + const { tmpPath } = await streamResumableDownload({ + url, + destPath, + headers, + fetchImpl, + finalize: false, + onHttpError, + onBytes: (received, total) => { + lastTick = { received, total }; + slot.track(received, total); + if (emitProgress) emitProgress(received, total); + }, + ...slot.downloadOptions(), + }); + // `finalize: false` means streamResumableDownload never had a chance to + // clean up its own etag sidecar (that only happens on ITS finalize path). + // Reaching here means the stream completed successfully — every branch + // below either moves or deletes tmpPath outright, never leaves it for a + // future resume, so the sidecar describing it is equally done. + await rmGuarded(etagPathFor(destPath), { force: true }).catch(() => {}); + if (onProgress) onProgress(lastTick); + // Atomic no-clobber finalize: `link` is POSIX-atomic and fails with EEXIST + // when destPath already exists (concurrent install that snuck past our + // pre-check). On success we unlink the tmp; on EEXIST we clean up and + // throw CIVITAI_ALREADY_INSTALLED. For other link errors (cross-device + // EXDEV, read-only fs, etc.) fall back to rename, which is the only + // portable option on those platforms. + const linkErr = await link(tmpPath, destPath).catch((e) => e); + if (!linkErr) { + await unlinkGuarded(tmpPath).catch(() => {}); + return; + } + if (linkErr.code === 'EEXIST') { + await rmGuarded(tmpPath, { force: true }).catch(() => {}); + const basename_ = basename(destPath); + throw new ServerError( + `Already installed: ${basename_}. Delete it first or pick a different version.`, + { status: 409, code: 'CIVITAI_ALREADY_INSTALLED' }, + ); + } + // EXDEV or similar — fall back to rename. Re-check destPath right before + // the rename so a concurrent install that landed between our link attempt + // and now can't be silently clobbered (POSIX rename overwrites). Treat + // late-arriving dest as CIVITAI_ALREADY_INSTALLED, matching the EEXIST + // path above. + if (existsSync(destPath)) { + await rmGuarded(tmpPath, { force: true }).catch(() => {}); + const basename_ = basename(destPath); + throw new ServerError( + `Already installed: ${basename_}. Delete it first or pick a different version.`, + { status: 409, code: 'CIVITAI_ALREADY_INSTALLED' }, + ); + } + await rename(tmpPath, destPath).catch(async (err) => { + await rmGuarded(tmpPath, { force: true }).catch(() => {}); + throw err; + }); + } catch (err) { + throw slot.wrapError(err); + } finally { + if (signal) signal.removeEventListener('abort', onExternalAbort); + slot.release(); } - await rename(tmpPath, destPath).catch(async (err) => { - await rmGuarded(tmpPath, { force: true }).catch(() => {}); - throw err; - }); }; // After a LoRA finishes downloading, verify the on-disk `.safetensors` before diff --git a/server/services/loras.test.js b/server/services/loras.test.js index 8591f8c01f..537d1f9173 100644 --- a/server/services/loras.test.js +++ b/server/services/loras.test.js @@ -988,25 +988,76 @@ describe('installFromHuggingface', () => { expect(ticks[ticks.length - 1]).toEqual({ received: 10, total: 0 }); }); - it('forwards an AbortSignal to the download fetch so an SSE disconnect can cancel it', async () => { + it('lets an external AbortSignal cancel the in-flight weights download so an SSE disconnect stops it', async () => { + // The download claims a slot keyed on destPath and downloads through the + // SLOT's own AbortController (so a stall/cancel gets the same typed error + // regardless of who triggered it) — an external `signal` (the route's + // res.on('close') controller) is wired to cancel THROUGH the slot rather + // than being handed to fetch directly. Prove the wiring by aborting mid- + // transfer and asserting the actual weights fetch's signal aborts too. const controller = new AbortController(); let sawSignal; const fetchImpl = async (url, opts) => { if (url.startsWith('https://huggingface.co/api/models/')) return mockJsonResponse(HF_MODEL); if (url.includes('/resolve/main/pytorch_lora_weights.safetensors')) { sawSignal = opts?.signal; - const stream = new ReadableStream({ start(c) { c.enqueue(new Uint8Array(validSafetensors())); c.close(); } }); - return { ok: true, status: 200, body: stream }; + // Never closes on its own: this stream only ends when the abort + // listener below errors it, mirroring how real fetch aborts an + // in-flight response body when its signal fires. + const stream = new ReadableStream({ + start(c) { + c.enqueue(new Uint8Array(validSafetensors().slice(0, 8))); + sawSignal.addEventListener('abort', () => c.error(new Error('aborted')), { once: true }); + }, + }); + return { ok: true, status: 200, body: stream, headers: new Map([['content-length', String(validSafetensors().length)]]) }; } throw new Error(`unexpected fetch: ${url}`); }; - await lorasService.installFromHuggingface( + const install = lorasService.installFromHuggingface( { url: 'https://huggingface.co/fal/ltx2.3-audio-reactive-lora', token: 'hf_test' }, { fetchImpl, signal: controller.signal }, ); - // The controller's signal must reach the actual weights download (not just - // the metadata fetch) — that's the transfer a disconnect needs to cancel. - expect(sawSignal).toBe(controller.signal); + await vi.waitFor(() => expect(sawSignal).toBeDefined()); + expect(sawSignal.aborted).toBe(false); + controller.abort(); + expect(sawSignal.aborted).toBe(true); + const err = await install.catch((e) => e); + expect(err.message).toMatch(/cancel/i); + expect(err.code).toBe('LORA_DOWNLOAD_CANCELLED'); + }); + + it('refuses a second install of the same destination while the first is still downloading', async () => { + let release; + const gate = new Promise((resolve) => { release = resolve; }); + const full = validSafetensors(); + const fetchImpl = async (url) => { + if (url.startsWith('https://huggingface.co/api/models/')) return mockJsonResponse(HF_MODEL); + if (url.includes('/resolve/main/pytorch_lora_weights.safetensors')) { + const stream = new ReadableStream({ + start(c) { + c.enqueue(new Uint8Array(full.slice(0, 8))); + gate.then(() => { c.enqueue(new Uint8Array(full.slice(8))); c.close(); }); + }, + }); + return { ok: true, status: 200, body: stream, headers: new Map([['content-length', String(full.length)]]) }; + } + throw new Error(`unexpected fetch: ${url}`); + }; + const first = lorasService.installFromHuggingface( + { url: 'https://huggingface.co/fal/ltx2.3-audio-reactive-lora', token: 'hf_test' }, + { fetchImpl }, + ); + const destPath = join(tmpLoras, 'lora-fal-ltx2.3-audio-reactive-lora-hf.safetensors'); + await vi.waitFor(() => expect(existsSync(`${destPath}.partial`)).toBe(true)); + const err = await lorasService.installFromHuggingface( + { url: 'https://huggingface.co/fal/ltx2.3-audio-reactive-lora', token: 'hf_test' }, + { fetchImpl }, + ).catch((e) => e); + expect(err).toBeInstanceOf(Error); + expect(err.code).toBe('LORA_DOWNLOAD_IN_FLIGHT'); + release(); + await first; }); it('installs a Flux.2 Klein 9B collection as flux2 and picks the klein9b file', async () => { @@ -1113,6 +1164,56 @@ describe('installFromHuggingface', () => { }); }); +// #6190: loras.js is the third `streamResumableDownload` caller to register a +// slot — the orphaned-partial GC sweeps PATHS.loras but, before this, couldn't +// see a live LoRA download's `.partial` because loras.js claimed no slot. +describe('loras.js download-slot registration with the orphaned-partial GC', () => { + const HF_MODEL = { + id: 'fal/ltx2.3-audio-reactive-lora', + tags: ['ltxv', 'lora'], + cardData: { base_model: 'Lightricks/LTX-2.3', instance_prompt: 'audio reactive' }, + siblings: [ + { rfilename: 'README.md' }, + { rfilename: 'pytorch_lora_weights.safetensors' }, + ], + }; + const ANCIENT = new Date(Date.now() - 30 * 24 * 60 * 60 * 1000); + + it('protects a live LoRA download from the sweep even once its .partial is aged past the gate', async () => { + const { sweepOrphanedDownloadPartials } = await import('./orphanedPartialGc.js'); + let release; + const gate = new Promise((resolve) => { release = resolve; }); + const fetchImpl = async (url) => { + if (url.startsWith('https://huggingface.co/api/models/')) return mockJsonResponse(HF_MODEL); + if (url.includes('/resolve/main/pytorch_lora_weights.safetensors')) { + const full = validSafetensors(); + const stream = new ReadableStream({ + start(c) { + c.enqueue(new Uint8Array(full.slice(0, 8))); + gate.then(() => { c.enqueue(new Uint8Array(full.slice(8))); c.close(); }); + }, + }); + return { ok: true, status: 200, body: stream, headers: new Map([['content-length', String(full.length)]]) }; + } + throw new Error(`unexpected fetch: ${url}`); + }; + const install = lorasService.installFromHuggingface( + { url: 'https://huggingface.co/fal/ltx2.3-audio-reactive-lora', token: 'hf_test' }, + { fetchImpl }, + ); + const partialPath = join(tmpLoras, 'lora-fal-ltx2.3-audio-reactive-lora-hf.safetensors.partial'); + await vi.waitFor(() => expect(existsSync(partialPath)).toBe(true)); + const fs = await import('fs/promises'); + await fs.utimes(partialPath, ANCIENT, ANCIENT); + + expect(await sweepOrphanedDownloadPartials({ dirs: [tmpLoras] })).toMatchObject({ deleted: 0, keptProtected: 1 }); + expect(existsSync(partialPath)).toBe(true); + + release(); + await install; + }); +}); + describe('LoRA key layout', () => { const writeLora = async (name, header) => { const fs = await import('fs/promises'); diff --git a/server/services/orphanedPartialGc.js b/server/services/orphanedPartialGc.js index 9f3b48c4fe..002ac530a6 100644 --- a/server/services/orphanedPartialGc.js +++ b/server/services/orphanedPartialGc.js @@ -20,11 +20,12 @@ import { resolveSpecModelPath } from './specDecodeModels.js'; import { getModelsDir as getOllamaModelsDir } from './ollamaManager.js'; import { getModelsDir as getLmStudioModelsDir } from './lmStudioManager.js'; import { slotstreamCacheDir } from '../lib/slotstreamModels.js'; -// Side-effect import: constructing the module is what registers its download +// Side-effect imports: constructing a module is what registers its download // slot with `isAnyDownloadInFlight`, and an unregistered slot's live shards are // unprotected from the sweep below. `orphanedPartialGc.test.js` fails if a // module calling `createDownloadSlot` stops being reachable from here. import './slotstreamModelManager.js'; +import './loras.js'; import { createSweepScheduler } from './sweepScheduler.js'; export { ORPHANED_PARTIAL_MAX_AGE_MS }; @@ -65,9 +66,7 @@ export async function sweepOrphanedDownloadPartials({ } = {}) { const targets = dirs || await collectPartialSweepDirs(); // One predicate for every runtime: each download slot registers itself, so a - // path that claims a slot is protected without a clause added here. (A path - // that claims none — `loras.js`, which writes into `PATHS.loras` above — is - // not; migrating it is issue #6190.) + // path that claims a slot is protected without a clause added here. return sweepOrphanedPartials(targets, { now, maxAgeMs, isProtected: isAnyDownloadInFlight }); }