diff --git a/docs/huggingface-model-import-plan.md b/docs/huggingface-model-import-plan.md index 7109cbc7aa..b874fc9253 100644 --- a/docs/huggingface-model-import-plan.md +++ b/docs/huggingface-model-import-plan.md @@ -1,6 +1,9 @@ # Importing models from Hugging Face -**Status:** built, not deployed. The migration has not been applied to any database. +**Status:** the page, the transfer and the *Manage files* picker are merged and live, and both +`20260914120000_huggingface_import` and `20260918120000_huggingface_import_attach_target` are applied +to production. The transfer API below is on `feat/huggingface-import-api`, not yet merged, and adds +`20260918140000_huggingface_import_pending_attach_idx`. ## The goal @@ -13,12 +16,14 @@ this as the first route. | Piece | Where | | --- | --- | -| `HuggingFaceImport` table + status enum | `packages/civitai-db-schema/prisma/schema.full.prisma`, migration `20260914120000_huggingface_import` | +| `HuggingFaceImport` table + status enum | `packages/civitai-db-schema/prisma/schema.full.prisma`, migrations `20260914120000_huggingface_import`, `20260918120000_huggingface_import_attach_target` (`attachVersionId`, `attachType`) and `20260918140000_huggingface_import_pending_attach_idx` | | HF API client — parse a URL, resolve a branch to a commit sha, list files with sizes and LFS sha256, ranged reads | `src/server/services/huggingface.service.ts` | | Queue + the resumable transfer | `src/server/services/huggingface-import.service.ts` | | The runner | `src/server/jobs/process-huggingface-imports.ts`, registered in the `jobs` array in `run-jobs` | -| tRPC surface (`getAll` — filterable by `groupName`/`repo`, and by `unattached` — `getCounts`, `lookup`, `enqueue`, `attach`, `detach`, `delete`, `renameGroup`, `retry`, `cancel`) | `src/server/routers/huggingface-import.router.ts` | +| tRPC surface (`getAll` — filterable by `groupName`/`repo`, and by `unattached` — `getCounts`, `getConfig`, `lookup`, `enqueue`, `attach`, `detach`, `delete`, `renameGroup`, `retry`, `cancel`, `setConfig`) | `src/server/routers/huggingface-import.router.ts` | | Moderator page | `src/pages/moderator/huggingface-import.tsx` + `src/components/Moderation/HuggingFaceImport/` | +| Transfer API — queue with a destination, poll one import, list a repo | `src/pages/api/admin/huggingface-import.ts` | +| The session-free create path the job attaches through | `createModelFile` in `src/server/controllers/model-file.controller.ts` | | The *Manage files* picker (moderator-only) | `src/components/Moderation/HuggingFaceImport/AddFromImportsModal.tsx`, opened via `src/components/Dialog/triggers/add-from-hugging-face-imports.ts` from `AddFromImportsButton.tsx` in `src/components/Resource/Files.tsx` | | Server-side multipart helpers (`createMultipartUpload`, `uploadPart`) | `src/utils/s3-utils.ts` | | The shared key builder (`buildUploadKey`) | `src/utils/upload-key.ts` — used by `/api/upload` **and** the import | @@ -62,14 +67,19 @@ bytes a model file still points at is refused either way. ## Attaching to a version -`attach` turns a finished import into a `ModelFile` on a version, going through `createFileHandler` -rather than writing the row directly — that is what gets the storage-resolver registration and the -inline scan submission, so scanning and hashing follow on their own. +`attach` turns a finished import into a `ModelFile` on a version, going through the shared create path in +`model-file.controller.ts` rather than writing the row directly — that is what gets the +storage-resolver registration and the inline scan submission, so scanning and hashing follow on their +own. The tRPC surfaces reach it as `createFileHandler`; the transfer job has no session to borrow, so +it calls `createModelFile` — the same body with `userId`, `isModerator` and `track` passed in. -Four ways in, one path underneath: the Attach control on a completed row, the **Add from Hugging +Five ways in, one path underneath: the Attach control on a completed row, the **Add from Hugging Face imports** picker inside a version's *Manage files*, the `huggingFaceImport.attach` procedure, -and two commands on `official-model-admin` — `hf-imports` (what transferred) and `attach-import` -(put one on a version). +`POST /api/admin/huggingface-import` (below), and two commands on `official-model-admin` — +`hf-imports` (what transferred) and `attach-import` (put one on a version). + +Four of those attach a file that has already landed. The API is the one that does not: it records +where each file is headed at enqueue, and the transfer job attaches it on completion. What that changes for the skill: the 20GB re-upload a person used to perform is gone, and attaching is scriptable. A human still queues the repo on the page, and still confirms the file type when the @@ -137,8 +147,9 @@ No real-world measurement exists yet — nothing has run against a live bucket. **Nothing checks that the file you attached is the file you meant.** The scan catches malware, not mislabelling. -The direction is a **pull**: an import usually happens before anyone knows which version will want -it, so the version draws from the pool rather than the import pushing at a version. Both surfaces +The default direction is a **pull**: an import usually happens before anyone knows which version +will want it, so the version draws from the pool rather than the import pushing at a version. The +transfer API is the one push — a caller that already knows the version has nothing to come back for. Both surfaces that make that work now exist — the **Unattached** tab on the import page and the **Add from Hugging Face imports** picker inside a version's *Manage files*, both grouped and filtered by `groupName`. @@ -151,10 +162,9 @@ understate what we hold. clears while leaving the `ModelFile` alive. Judging from the row alone destroys the bytes a published version is serving, two clicks after a detach. -**No quota.** Moderator-only at the router; everything underneath is already scoped per owner, so -opening it up needs a per-user quota — size and count — and a `userId` in the -`(repo, revision, filename)` unique index — the header comment in `huggingface-import.router.ts` is -the prerequisite list. +**No quota.** Moderator-only at the router; everything underneath is already scoped per owner. What +opening it up needs is listed under *The transfer API* below, and in the header comment of +`huggingface-import.router.ts`. **Licenses are shown, not enforced.** The page surfaces the declared license and flags gated repos. @@ -177,6 +187,45 @@ The `Import` table, its `ImportStatus` enum, and the `fromImportId` columns on ` `ModelVersion` are left in place: dropping them is a migration over existing rows, and nothing reads them any more. +## The transfer API + +`POST /api/admin/huggingface-import` is internal tooling only — the same operational surface as the +rest of `src/pages/api/admin/`, and not something a creator's own tooling reaches. + +``` +POST /api/admin/huggingface-import?token=… +{ "repo": "black-forest-labs/FLUX.1-dev", "modelVersionId": 123, + "files": [{ "path": "flux1-dev.safetensors", "type": "Model" }, + { "path": "ae.safetensors", "type": "VAE" }] } +→ { repo, revision, files: [{ path, importId, status, modelFileId }] } + +GET /api/admin/huggingface-import?token=…&id=11 # one import's progress +GET /api/admin/huggingface-import?token=…&repo=… # what is in the repo, with sizes and hashes +``` + +The caller polls `id` and does nothing else — `attachVersionId` and `attachType` are recorded on the +row, and `attachIfRequested` runs when the bytes land. A failed attach leaves the file transferred and +unattached with the reason on the row, where the **Unattached** tab shows it; it is not a failed +transfer, and the row does not say it is. + +Type is required per file and refused when the extension cannot be it, which is the same rule the +upload UI applies — checked before a transfer that can run for hours, not after. So is the target +version: an id nothing matches is a 400, not an hours-late attach failure. + +**A file whose sha256 we already store is attached without transferring anything** (`reused: true` in +the response, and nothing to poll). Hugging Face publishes each LFS file's sha in its tree, so this is +decided before any bytes move — a text encoder shared by a dozen repos costs one `ModelFileHash` +lookup we already run. The match is on the sha and never on the filename: `ae.safetensors` and +`model.safetensors` name different bytes in different repos, and the wrong weights on a version are +invisible until someone generates. The import row records the provenance and carries no `bucket`/`key` +— the object belongs to the file that first stored it, and the refcounted delete already refuses to +free bytes another `ModelFile` still points at. + +**A user-facing version is a different surface, not a flag on this one.** It needs everything the +router's header lists — a per-user quota, a `userId` in the `(repo, revision, filename)` unique index +— plus the target version checked against the caller, and gated repos refused: the importer +authenticates as one shared account, and its accepted terms are ours rather than theirs. + ## Open questions 1. **Who may import**, and under what quota — see above. diff --git a/docs/moderator-app/page-migration-checklist.md b/docs/moderator-app/page-migration-checklist.md index 902d2cdc7f..59170ebb45 100644 --- a/docs/moderator-app/page-migration-checklist.md +++ b/docs/moderator-app/page-migration-checklist.md @@ -494,7 +494,7 @@ Real, working features the head moderator doesn't use or doesn't know about. Mig - Services: `huggingface.service.ts` (HF API client), `huggingface-import.service.ts` (queue + resumable multipart transfer), `huggingface-import-config.service.ts` (transfer settings) - Schemas: `huggingface-import.schema.ts` - Infra: **Postgres (`HuggingFaceImport`) + S3/B2 multipart + Redis (sysRedis transfer config) + the `process-huggingface-imports` cron** - - Notes: a port moves the page, not the transfer — the cron and `createFileHandler` (scan + hash submission) stay in the main app, so this is delegate-shaped if it moves at all. + - Notes: a port moves the page, not the transfer — the cron and the model-file create path (`createFileHandler` for the tRPC surfaces, `createModelFile` for the cron's auto-attach; both carry the scan + hash submission) stay in the main app, so this is delegate-shaped if it moves at all. --- diff --git a/packages/civitai-db-schema/prisma/migrations/20260918120000_huggingface_import_attach_target/migration.sql b/packages/civitai-db-schema/prisma/migrations/20260918120000_huggingface_import_attach_target/migration.sql new file mode 100644 index 0000000000..1ae7264bbc --- /dev/null +++ b/packages/civitai-db-schema/prisma/migrations/20260918120000_huggingface_import_attach_target/migration.sql @@ -0,0 +1,5 @@ +-- Where a queued import is headed, so the transfer job can attach the file when it lands. +-- Both nullable with no default: existing rows are imports nobody asked to auto-attach. +ALTER TABLE "HuggingFaceImport" + ADD COLUMN "attachVersionId" INTEGER, + ADD COLUMN "attachType" TEXT; diff --git a/packages/civitai-db-schema/prisma/migrations/20260918140000_huggingface_import_pending_attach_idx/migration.sql b/packages/civitai-db-schema/prisma/migrations/20260918140000_huggingface_import_pending_attach_idx/migration.sql new file mode 100644 index 0000000000..58cd9b759d --- /dev/null +++ b/packages/civitai-db-schema/prisma/migrations/20260918140000_huggingface_import_pending_attach_idx/migration.sql @@ -0,0 +1,16 @@ +-- The attach sweep runs every cron tick and orders by "completedAt", which no index covers: without +-- this it heap-reads every Completed row, filters, sorts, and only then takes its one row — on the +-- primary, and in the steady state to find nothing. `Completed` is terminal and nothing prunes it, +-- so the scanned set only grows. +-- +-- The predicate matches the query's WHERE term for term, which is what lets the planner use it. A +-- failed attach records `error`, so a poisoned row drops straight out of the index and it stays +-- near-empty rather than tracking the table. +-- +-- CONCURRENTLY cannot run inside a transaction block — run this statement on its own. +CREATE INDEX CONCURRENTLY IF NOT EXISTS "HuggingFaceImport_pending_attach_idx" + ON "HuggingFaceImport" ("completedAt") + WHERE status = 'Completed' + AND "modelFileId" IS NULL + AND "attachVersionId" IS NOT NULL + AND error IS NULL; diff --git a/packages/civitai-db-schema/prisma/schema.full.prisma b/packages/civitai-db-schema/prisma/schema.full.prisma index ed4ad9b018..eed9524a43 100644 --- a/packages/civitai-db-schema/prisma/schema.full.prisma +++ b/packages/civitai-db-schema/prisma/schema.full.prisma @@ -893,6 +893,10 @@ model HuggingFaceImport { userId Int? modelVersionId Int? modelFileId Int? + /// Where this file is headed, recorded before the bytes move so the transfer job can attach it + /// itself. `modelVersionId` cannot carry this: it means "attached to", and detaching clears it. + attachVersionId Int? + attachType String? /// Worker lease. A transfer outlives any one job run, so a claim plus a heartbeat is what stops two /// runs moving the same file and what lets the next run tell "in flight" from "abandoned". claimedBy String? @@ -906,6 +910,10 @@ model HuggingFaceImport { @@unique([repo, revision, filename]) @@index([status, createdAt]) @@index([groupName]) + /// Also `HuggingFaceImport_pending_attach_idx` — a PARTIAL index on ("completedAt") for the attach + /// sweep, which Prisma cannot express, so it lives in its migration and will not round-trip here. + /// Also `HuggingFaceImport_pending_attach_idx` — a PARTIAL index on ("completedAt") for the attach + /// sweep, which Prisma cannot express, so it lives in its migration and will not round-trip here. } enum ModelStatus { diff --git a/packages/civitai-db-schema/src/kysely/types.ts b/packages/civitai-db-schema/src/kysely/types.ts index 5ee7cdbefe..7282f5b6d8 100644 --- a/packages/civitai-db-schema/src/kysely/types.ts +++ b/packages/civitai-db-schema/src/kysely/types.ts @@ -2539,6 +2539,12 @@ export type HuggingFaceImport = { userId: number | null; modelVersionId: number | null; modelFileId: number | null; + /** + * Where this file is headed, recorded before the bytes move so the transfer job can attach it + * itself. `modelVersionId` cannot carry this: it means "attached to", and detaching clears it. + */ + attachVersionId: number | null; + attachType: string | null; /** * Worker lease. A transfer outlives any one job run, so a claim plus a heartbeat is what stops two * runs moving the same file and what lets the next run tell "in flight" from "abandoned". diff --git a/packages/civitai-db-schema/src/models.ts b/packages/civitai-db-schema/src/models.ts index 3fe612ed63..06b01718fa 100644 --- a/packages/civitai-db-schema/src/models.ts +++ b/packages/civitai-db-schema/src/models.ts @@ -862,6 +862,8 @@ export interface HuggingFaceImport { userId: number | null; modelVersionId: number | null; modelFileId: number | null; + attachVersionId: number | null; + attachType: string | null; claimedBy: string | null; claimedAt: Date | null; heartbeatAt: Date | null; diff --git a/src/pages/api/admin/huggingface-import.ts b/src/pages/api/admin/huggingface-import.ts new file mode 100644 index 0000000000..96bc3ead7e --- /dev/null +++ b/src/pages/api/admin/huggingface-import.ts @@ -0,0 +1,199 @@ +import type { NextApiRequest, NextApiResponse } from 'next'; +import * as z from 'zod'; +import { dbRead } from '~/server/db/client'; +import { + transferHuggingFaceImportSchema, + getHuggingFaceImportsSchema, +} from '~/server/schema/huggingface-import.schema'; +import { HuggingFaceError, parseHuggingFaceRepo } from '~/server/services/huggingface.service'; +import { + enqueueImports, + getImportStatus, + IMPORT_SYSTEM_USER_ID, + resolveRepoForImport, + reuseStoredFile, +} from '~/server/services/huggingface-import.service'; +import { WebhookEndpoint } from '~/server/utils/endpoint-helpers'; +import { filterFileTypeByExtension } from '~/utils/file-display-helpers'; +import { zc } from '~/utils/schema-helpers'; + +/** + * Server-side Hugging Face import, for our own tooling. + * + * POST — queue a repo's files onto a model version; one call can carry the whole model. + * { repo, revision?, modelVersionId, groupName?, userId?, files: [{ path, type }] } + * `userId` is attribution only and defaults to the system user. + * GET `?id=` — one import's progress. + * GET `?repo=&revision=` — the repo's files with sizes, sha256, and whether we already store the sha. + * + * The transfer runs on the `process-huggingface-imports` cron and attaches each file when it lands, + * so a caller polls `id` and makes no second call. A file whose sha256 we already store is attached + * immediately instead — `reused: true`, with nothing to poll. Response shapes are whatever + * `getImportStatus` and `resolveRepoForImport` select. + * + * 🔴 `type` is required per file and never inferred on POST. The GET's `suggestedType` is a hint for a + * human and never names the primary weights, because a wrong type there is a version nothing loads. + */ +const getStatusSchema = z.object({ id: zc.numberString }); +const getRepoSchema = getHuggingFaceImportsSchema + .pick({ repo: true }) + .required({ repo: true }) + .extend({ revision: z.string().trim().min(1).optional() }); + +/** Hugging Face's own refusals are the caller's problem to fix, not a 500. */ +async function resolveRepo( + input: { repo: string; revision?: string }, + res: NextApiResponse +): Promise> | null> { + try { + // Parsed first, exactly as the tRPC surface does it: a pasted repo URL — or a `/tree/` one + // whose revision must be kept — is interpolated raw into the HF path otherwise, and comes back + // as a 404 that reads like a missing repo. + const target = parseHuggingFaceRepo(input.repo); + if (!target) { + res.status(400).json({ error: 'Could not read an owner/name out of that.' }); + return null; + } + // An explicit revision wins; otherwise a tree URL keeps the revision it names rather than + // silently resolving to the default branch. + return await resolveRepoForImport({ + repo: target.repo, + revision: input.revision ?? target.revision, + }); + } catch (error) { + if (error instanceof HuggingFaceError) { + res.status(400).json({ error: error.message }); + return null; + } + throw error; + } +} + +export default WebhookEndpoint(async (req: NextApiRequest, res: NextApiResponse) => { + if (req.method !== 'GET' && req.method !== 'POST') { + res.setHeader('Allow', ['GET', 'POST']); + return res.status(405).json({ error: `${req.method} not allowed` }); + } + + if (req.method === 'GET') { + if (req.query.repo) { + const query = getRepoSchema.safeParse(req.query); + if (!query.success) return res.status(400).json({ error: z.prettifyError(query.error) }); + const listed = await resolveRepo(query.data, res); + if (!listed) return; + return res.status(200).json(listed); + } + + const query = getStatusSchema.safeParse(req.query); + if (!query.success) return res.status(400).json({ error: z.prettifyError(query.error) }); + const status = await getImportStatus(query.data.id); + if (!status) return res.status(404).json({ error: `No import with id ${query.data.id}.` }); + return res.status(200).json(status); + } + + // `safeParse`, not `parse`: a thrown ZodError leaves this handler as an uncaught 500, and a + // malformed body is the caller's mistake to read. + const parsed = transferHuggingFaceImportSchema.safeParse(req.body); + if (!parsed.success) return res.status(400).json({ error: z.prettifyError(parsed.error) }); + const input = parsed.data; + + // Before the transfer, not after: an unknown version id otherwise costs hours of copying and then + // fails at the attach, with the bytes already stored. + const version = await dbRead.modelVersion.findUnique({ + where: { id: input.modelVersionId }, + select: { id: true }, + }); + if (!version) + return res.status(400).json({ error: `No model version with id ${input.modelVersionId}.` }); + + const repo = await resolveRepo(input, res); + if (!repo) return; + + const available = new Set(repo.files.map((file) => file.path)); + const missing = input.files.filter((file) => !available.has(file.path)).map((file) => file.path); + if (missing.length) + return res.status(400).json({ + error: `Not in ${repo.repo} at ${repo.revision}: ${missing.join(', ')}`, + }); + + // The rule the upload UI applies, before a transfer that can run for hours. + const mistyped = input.files.filter((file) => !filterFileTypeByExtension(file.type, file.path)); + if (mistyped.length) + return res.status(400).json({ + error: mistyped.map((file) => `${file.path} cannot be "${file.type}"`).join('; '), + }); + + const userId = input.userId ?? IMPORT_SYSTEM_USER_ID; + const inRepo = new Map(repo.files.map((file) => [file.path, file])); + + // Split on the sha Hugging Face already gave us: bytes we hold are attached now, the rest are + // queued. A shared 10 GB text encoder is the difference between a transfer and nothing at all. + const reused = new Map(); + for (const file of input.files) { + const source = inRepo.get(file.path); + // `existing` is a sha256 match, never a filename one: the same name carries different bytes in + // different repos, and the wrong weights on a version are invisible until someone generates. + if (!source?.existing) continue; + reused.set( + file.path, + await reuseStoredFile({ + repo: repo.repo, + revision: repo.revision, + path: file.path, + sizeBytes: source.size ?? null, + sha256: source.sha256 ?? null, + storedUrl: source.existing.url, + modelVersionId: input.modelVersionId, + type: file.type, + userId, + groupName: input.groupName, + }) + ); + } + + const toTransfer = input.files.filter((file) => !reused.has(file.path)); + const { rows } = toTransfer.length + ? await enqueueImports({ + repo, + paths: toTransfer.map((file) => file.path), + userId, + groupName: input.groupName, + attach: { + modelVersionId: input.modelVersionId, + types: Object.fromEntries(toTransfer.map((file) => [file.path, file.type])), + }, + }) + : { rows: [] }; + + const byPath = new Map(rows.map((row) => [row.filename, row])); + return res.status(200).json({ + repo: repo.repo, + revision: repo.revision, + files: input.files.map((file) => { + const alreadyStored = reused.get(file.path); + if (alreadyStored) + return { + path: file.path, + importId: alreadyStored.importId, + status: 'Completed', + modelFileId: alreadyStored.modelFileId, + attachVersionId: input.modelVersionId, + reused: true, + }; + + const row = byPath.get(file.path); + return { + path: file.path, + importId: row?.id ?? null, + status: row?.status ?? 'Unknown', + modelFileId: row?.modelFileId ?? null, + reused: false, + // 🔴 Where the file is ACTUALLY headed. `(repo, revision, filename)` is unique, so a path + // queued earlier keeps its original row and its original destination — this call did not + // re-point it, and without this field a caller polling to `Completed` would conclude its + // own version got the file. + attachVersionId: row?.attachVersionId ?? null, + }; + }), + }); +}); diff --git a/src/server/__tests__/huggingface-import-endpoint.test.ts b/src/server/__tests__/huggingface-import-endpoint.test.ts new file mode 100644 index 0000000000..3f97532e55 --- /dev/null +++ b/src/server/__tests__/huggingface-import-endpoint.test.ts @@ -0,0 +1,310 @@ +import { beforeEach, describe, expect, it, vi } from 'vitest'; +import type * as HuggingFaceImportService from '~/server/services/huggingface-import.service'; +// Side-effect imports: the canonical mocks the handler's graph reaches at import time, WEBHOOK_TOKEN +// among them. +import '~/__tests__/mocks/logging.mock'; +import { dbMock } from '~/__tests__/mocks/db.mock'; +import '~/__tests__/mocks/env.mock'; + +const { resolveRepoForImport, enqueueImports, getImportStatus, reuseStoredFile } = vi.hoisted( + () => ({ + resolveRepoForImport: vi.fn(), + enqueueImports: vi.fn(), + getImportStatus: vi.fn(), + reuseStoredFile: vi.fn(), + }) +); + +vi.mock('~/server/services/huggingface-import.service', async (importOriginal) => ({ + ...(await importOriginal()), + resolveRepoForImport, + enqueueImports, + getImportStatus, + reuseStoredFile, +})); + +import { IMPORT_SYSTEM_USER_ID } from '~/server/services/huggingface-import.service'; + +const handler = (await import('~/pages/api/admin/huggingface-import')).default; + +const dbRead = dbMock.dbRead; + +function run({ + method = 'POST', + query = {}, + body, + token = 'test-webhook-token', +}: { + method?: string; + query?: Record; + body?: unknown; + token?: string | null; +}) { + const req = { + method, + query: { ...(token ? { token } : {}), ...query }, + headers: {}, + body, + }; + + let statusCode = 0; + let payload: unknown; + const res = { + status(code: number) { + statusCode = code; + return res; + }, + json(data: unknown) { + payload = data; + return res; + }, + send: () => res, + setHeader: () => res, + end: () => res, + _status: () => statusCode, + _body: () => payload as Record, + }; + + return handler(req as never, res as never).then(() => res); +} + +const repo = { + repo: 'black-forest-labs/FLUX.1-dev', + revision: 'abc123', + files: [ + { path: 'flux1-dev.safetensors', size: 100, sha256: 'a', suggestedType: null, existing: null }, + { path: 'ae.safetensors', size: 50, sha256: 'b', suggestedType: 'VAE', existing: null }, + ], +}; + +const post = (overrides: Record = {}) => ({ + repo: 'black-forest-labs/FLUX.1-dev', + modelVersionId: 42, + files: [ + { path: 'flux1-dev.safetensors', type: 'Model' }, + { path: 'ae.safetensors', type: 'VAE' }, + ], + ...overrides, +}); + +const queuedRows = [ + { + id: 11, + filename: 'flux1-dev.safetensors', + status: 'Queued', + modelFileId: null, + attachVersionId: 42, + }, + { id: 12, filename: 'ae.safetensors', status: 'Queued', modelFileId: null, attachVersionId: 42 }, +]; + +beforeEach(() => { + vi.clearAllMocks(); + resolveRepoForImport.mockResolvedValue(repo); + enqueueImports.mockResolvedValue({ queued: 2, skipped: 0, rows: queuedRows }); + dbRead.modelVersion.findUnique.mockResolvedValue({ id: 42 }); + reuseStoredFile.mockResolvedValue({ importId: 5, modelFileId: 900 }); +}); + +describe('POST /api/admin/huggingface-import', () => { + it('queues each file with the type it was given, and returns ids to poll', async () => { + const res = await run({ body: post({ groupName: 'flux', userId: 7 }) }); + + expect(res._status()).toBe(200); + expect(enqueueImports).toHaveBeenCalledTimes(1); + expect(enqueueImports).toHaveBeenCalledWith({ + repo, + paths: ['flux1-dev.safetensors', 'ae.safetensors'], + userId: 7, + groupName: 'flux', + attach: { + modelVersionId: 42, + types: { 'flux1-dev.safetensors': 'Model', 'ae.safetensors': 'VAE' }, + }, + }); + expect(res._body()).toMatchObject({ + revision: 'abc123', + files: [ + { path: 'flux1-dev.safetensors', importId: 11, status: 'Queued', attachVersionId: 42 }, + { path: 'ae.safetensors', importId: 12, status: 'Queued', attachVersionId: 42 }, + ], + }); + }); + + it('files a tooling import under the system user when no userId is given', async () => { + await run({ body: post() }); + + expect(enqueueImports.mock.calls[0][0].userId).toBe(IMPORT_SYSTEM_USER_ID); + }); + + it('says where a re-queued path is ACTUALLY headed', async () => { + // `(repo, revision, filename)` is unique, so a path queued earlier keeps its original row and + // its original destination. Without this the caller polls to `Completed` and concludes its own + // version got the file. + enqueueImports.mockResolvedValue({ + queued: 0, + skipped: 2, + rows: [ + { ...queuedRows[0], status: 'Completed', modelFileId: 900, attachVersionId: 7 }, + queuedRows[1], + ], + }); + + const res = await run({ body: post() }); + + expect(res._body().files).toMatchObject([ + { path: 'flux1-dev.safetensors', status: 'Completed', modelFileId: 900, attachVersionId: 7 }, + { path: 'ae.safetensors', status: 'Queued', attachVersionId: 42 }, + ]); + }); + + it('accepts a pasted repo URL, keeping the revision it names', async () => { + // Interpolated raw into the Hugging Face path otherwise, which comes back as a 404 that reads + // like a missing repo. + await run({ + body: post({ repo: 'https://huggingface.co/black-forest-labs/FLUX.1-dev/tree/beef123' }), + }); + + expect(resolveRepoForImport).toHaveBeenCalledWith({ + repo: 'black-forest-labs/FLUX.1-dev', + revision: 'beef123', + }); + }); + + it('refuses a version that does not exist, before queueing anything', async () => { + dbRead.modelVersion.findUnique.mockResolvedValue(null); + + const res = await run({ body: post({ modelVersionId: 999 }) }); + + expect(res._status()).toBe(400); + expect(String(res._body().error)).toContain('999'); + expect(enqueueImports).not.toHaveBeenCalled(); + }); + + it('refuses a file the repo does not have, before queueing anything', async () => { + const res = await run({ + body: post({ files: [{ path: 'not-there.safetensors', type: 'Model' }] }), + }); + + expect(res._status()).toBe(400); + expect(String(res._body().error)).toContain('not-there.safetensors'); + expect(enqueueImports).not.toHaveBeenCalled(); + }); + + it('refuses a type the extension cannot be', async () => { + const res = await run({ + body: post({ files: [{ path: 'flux1-dev.safetensors', type: 'Config' }] }), + }); + + expect(res._status()).toBe(400); + expect(String(res._body().error)).toContain('flux1-dev.safetensors'); + expect(enqueueImports).not.toHaveBeenCalled(); + }); + + it('answers a malformed body with a 400, not an uncaught 500', async () => { + const res = await run({ body: post({ files: [{ path: 'flux1-dev.safetensors' }] }) }); + + expect(res._status()).toBe(400); + expect(enqueueImports).not.toHaveBeenCalled(); + }); + + it('attaches a file we already store instead of transferring it again', async () => { + // Hugging Face publishes the sha before any bytes move, so this costs nothing to decide. + resolveRepoForImport.mockResolvedValue({ + ...repo, + files: [ + repo.files[0], + { + ...repo.files[1], + existing: { fileId: 77, name: 'ae.safetensors', url: 'https://s3.example/x' }, + }, + ], + }); + enqueueImports.mockResolvedValue({ queued: 1, skipped: 0, rows: [queuedRows[0]] }); + + const res = await run({ body: post() }); + + // Only the file we do not hold is queued. + expect(enqueueImports.mock.calls[0][0].paths).toEqual(['flux1-dev.safetensors']); + expect(reuseStoredFile).toHaveBeenCalledTimes(1); + expect(reuseStoredFile.mock.calls[0][0]).toMatchObject({ + path: 'ae.safetensors', + type: 'VAE', + storedUrl: 'https://s3.example/x', + modelVersionId: 42, + }); + expect(res._body().files).toMatchObject([ + { path: 'flux1-dev.safetensors', reused: false, importId: 11 }, + { path: 'ae.safetensors', reused: true, modelFileId: 900, status: 'Completed' }, + ]); + }); + + it('queues nothing at all when every file is one we hold', async () => { + resolveRepoForImport.mockResolvedValue({ + ...repo, + files: repo.files.map((file) => ({ + ...file, + existing: { fileId: 77, name: file.path, url: 'https://s3.example/x' }, + })), + }); + + const res = await run({ body: post() }); + + expect(enqueueImports).not.toHaveBeenCalled(); + expect(res._status()).toBe(200); + expect(res._body().files).toMatchObject([{ reused: true }, { reused: true }]); + }); + + it("passes Hugging Face's own refusal back as a 400", async () => { + const { HuggingFaceError } = await import('~/server/services/huggingface.service'); + resolveRepoForImport.mockRejectedValue(new HuggingFaceError('The repo is gated or private.')); + + const res = await run({ body: post() }); + + expect(res._status()).toBe(400); + expect(res._body().error).toBe('The repo is gated or private.'); + }); +}); + +describe('GET /api/admin/huggingface-import', () => { + it('reports one import by id', async () => { + getImportStatus.mockResolvedValue({ id: 11, status: 'Transferring', bytesTransferred: 10 }); + + const res = await run({ method: 'GET', query: { id: '11' } }); + + expect(getImportStatus).toHaveBeenCalledWith(11); + expect(res._body()).toMatchObject({ id: 11, status: 'Transferring' }); + }); + + it('404s an id that does not exist', async () => { + getImportStatus.mockResolvedValue(null); + + const res = await run({ method: 'GET', query: { id: '11' } }); + + expect(res._status()).toBe(404); + }); + + it('lists what is in a repo, so a caller can find the paths', async () => { + const res = await run({ method: 'GET', query: { repo: 'black-forest-labs/FLUX.1-dev' } }); + + expect(res._status()).toBe(200); + expect(res._body()).toMatchObject({ revision: 'abc123' }); + expect(getImportStatus).not.toHaveBeenCalled(); + }); +}); + +describe('the endpoint itself', () => { + it('refuses a method it does not implement', async () => { + const res = await run({ method: 'DELETE' }); + + expect(res._status()).toBe(405); + expect(enqueueImports).not.toHaveBeenCalled(); + }); + + it('refuses a caller without the token', async () => { + const res = await run({ method: 'GET', query: { id: '11' }, token: null }); + + expect(res._status()).toBe(401); + expect(getImportStatus).not.toHaveBeenCalled(); + }); +}); diff --git a/src/server/controllers/__tests__/model-file.create-context.test.ts b/src/server/controllers/__tests__/model-file.create-context.test.ts new file mode 100644 index 0000000000..0df7e1b2c4 --- /dev/null +++ b/src/server/controllers/__tests__/model-file.create-context.test.ts @@ -0,0 +1,77 @@ +import { beforeEach, describe, expect, it, vi } from 'vitest'; +import type * as ModelFileService from '~/server/services/model-file.service'; +import '~/__tests__/mocks/logging.mock'; +import '~/__tests__/mocks/db.mock'; +import '~/__tests__/mocks/env.mock'; + +/** + * `isModerator` is what lets `createFile` attach to a version the caller does not own. The transfer + * job passes it as a plain `true`, so the tRPC path must keep passing the session's own value — + * these pin that the two callers are not interchangeable. + */ +const { createFile, registerFileLocation, createModelFileScanRequest } = vi.hoisted(() => ({ + createFile: vi.fn(), + registerFileLocation: vi.fn(), + createModelFileScanRequest: vi.fn(), +})); + +vi.mock('~/server/services/model-file.service', async (importOriginal) => ({ + ...(await importOriginal()), + createFile, +})); +vi.mock('~/utils/storage-resolver', () => ({ registerFileLocation })); +vi.mock('~/server/services/model-file-scan.service', () => ({ + createModelFileScanRequest, + ModelFileScanSubmissionError: class extends Error {}, +})); + +import { createFileHandler, createModelFile } from '~/server/controllers/model-file.controller'; + +const input = { + modelVersionId: 42, + type: 'Model' as const, + name: 'flux.safetensors', + url: 'https://s3.example/model-bucket/model/7/flux.safetensors', + sizeKB: 2, +}; + +const track = { modelFile: vi.fn().mockResolvedValue(undefined) }; + +beforeEach(() => { + vi.clearAllMocks(); + createFile.mockResolvedValue({ + id: 900, + name: 'flux.safetensors', + modelVersion: { + id: 42, + modelId: 5, + baseModel: 'Flux.1 D', + status: 'Draft', + model: { type: 'Checkpoint' }, + _count: { posts: 0 }, + }, + }); + createModelFileScanRequest.mockResolvedValue(undefined); +}); + +describe('the model-file create path', () => { + it.each([true, false])( + "carries the session's isModerator (%s) through the tRPC path", + async (isModerator) => { + await createFileHandler({ + input, + ctx: { user: { id: 7, isModerator }, track } as never, + }); + + expect(createFile).toHaveBeenCalledWith(expect.objectContaining({ userId: 7, isModerator })); + } + ); + + it('lets a caller with no session state who it is acting as', async () => { + await createModelFile({ input, userId: -1, isModerator: true, track: track as never }); + + expect(createFile).toHaveBeenCalledWith( + expect.objectContaining({ userId: -1, isModerator: true }) + ); + }); +}); diff --git a/src/server/controllers/model-file.controller.ts b/src/server/controllers/model-file.controller.ts index 2f08444786..b951c4258a 100644 --- a/src/server/controllers/model-file.controller.ts +++ b/src/server/controllers/model-file.controller.ts @@ -1,5 +1,6 @@ import { TRPCError } from '@trpc/server'; import type { ProtectedContext } from '~/server/createContext'; +import type { Tracker } from '~/server/clickhouse/client'; import type { GetByIdInput } from '~/server/schema/base.schema'; import type { HasOfficialFileOfSizeInput, @@ -133,6 +134,28 @@ export const createFileHandler = async ({ }: { input: ModelFileCreateInput; ctx: ProtectedContext; +}) => + createModelFile({ + input, + userId: ctx.user.id, + isModerator: !!ctx.user.isModerator, + track: ctx.track, + }); + +/** + * The same create path with no request behind it, so a job can run it: a Hugging Face import + * attaches its own file when the transfer lands, and has no session to borrow. + */ +export const createModelFile = async ({ + input, + userId, + isModerator, + track, +}: { + input: ModelFileCreateInput; + userId: number; + isModerator: boolean; + track: Tracker; }) => { try { // Extract B2-specific fields before passing to createFile (they aren't DB columns) @@ -140,8 +163,8 @@ export const createFileHandler = async ({ const file = await createFile({ ...createInput, - userId: ctx.user.id, - isModerator: ctx.user.isModerator, + userId: userId, + isModerator: isModerator, select: { id: true, name: true, @@ -165,7 +188,7 @@ export const createFileHandler = async ({ if (resolved) { await safeRegisterFileLocation({ op: 'create', - userId: ctx.user.id, + userId: userId, fileId: file.id, modelVersionId: file.modelVersion.id, modelId: file.modelVersion.modelId, @@ -175,7 +198,7 @@ export const createFileHandler = async ({ }); } - ctx.track + track .modelFile({ type: 'Create', id: file.id, modelVersionId: file.modelVersion.id }) .catch(handleLogError); diff --git a/src/server/schema/huggingface-import.schema.ts b/src/server/schema/huggingface-import.schema.ts index 4cca81c72b..6c2334f26c 100644 --- a/src/server/schema/huggingface-import.schema.ts +++ b/src/server/schema/huggingface-import.schema.ts @@ -56,5 +56,28 @@ export const renameHuggingFaceGroupSchema = z.object({ groupName: z.string().trim().min(1).max(120), }); +/** `POST /api/admin/huggingface-import`. Composed from the enqueue shape so the two cannot drift. */ +export type TransferHuggingFaceImportInput = z.infer; +export const transferHuggingFaceImportSchema = enqueueHuggingFaceImportSchema + .omit({ paths: true }) + .extend({ + /** Optional here, unlike the UI enqueue: a caller may name a branch, tag or sha, or let the + * repo default stand. */ + revision: z.string().trim().min(1).optional(), + modelVersionId: z.number().int().positive(), + /** Attribution only. Defaults to the system user, which is what tooling imports are filed under. */ + userId: z.number().int().positive().optional(), + files: z + .array( + z.object({ + path: z.string().trim().min(1), + // Explicit, never inferred: this is what decides whether the version is loadable. + type: z.enum(constants.modelFileTypes), + }) + ) + .min(1) + .max(100), + }); + export type SetHuggingFaceImportConfigInput = z.infer; export const setHuggingFaceImportConfigSchema = huggingFaceImportConfigSchema.partial(); diff --git a/src/server/services/__tests__/huggingface-import.service.test.ts b/src/server/services/__tests__/huggingface-import.service.test.ts index 4abcdaa1d9..5c10b1e83d 100644 --- a/src/server/services/__tests__/huggingface-import.service.test.ts +++ b/src/server/services/__tests__/huggingface-import.service.test.ts @@ -2,6 +2,7 @@ import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest'; import { dbMock } from '~/__tests__/mocks/db.mock'; import { redisMock } from '~/__tests__/mocks/redis.mock'; import type * as HuggingFaceService from '~/server/services/huggingface.service'; +import type * as ModelFileController from '~/server/controllers/model-file.controller'; import type * as S3Utils from '~/utils/s3-utils'; import { setEnv } from '~/__tests__/mocks/env.mock'; @@ -13,6 +14,8 @@ const { mockComplete, mockAbort, mockDeleteObject, + mockObjectExists, + mockCreateModelFile, mockUrlsSafeToDelete, } = vi.hoisted(() => ({ mockReadRange: vi.fn(), @@ -22,6 +25,8 @@ const { mockComplete: vi.fn(), mockAbort: vi.fn(), mockDeleteObject: vi.fn(), + mockObjectExists: vi.fn(), + mockCreateModelFile: vi.fn(), mockUrlsSafeToDelete: vi.fn(), })); @@ -52,9 +57,15 @@ vi.mock('~/utils/s3-utils', async (importOriginal) => ({ completeMultipartUpload: mockComplete, abortMultipartUpload: mockAbort, deleteObject: mockDeleteObject, + objectExists: mockObjectExists, urlsSafeToDelete: mockUrlsSafeToDelete, })); +vi.mock('~/server/controllers/model-file.controller', async (importOriginal) => ({ + ...(await importOriginal()), + createModelFile: mockCreateModelFile, +})); + import { parseHuggingFaceRepo, suggestFileType } from '~/server/services/huggingface.service'; import { getHuggingFaceImportConfig, @@ -69,6 +80,12 @@ import { processImportQueue, renameGroup, retryImport, + attachIfRequested, + detachImport, + enqueueImports, + getImportStatus, + IMPORT_SYSTEM_USER_ID, + reuseStoredFile, } from '~/server/services/huggingface-import.service'; /** The width under test. Passed in rather than read from config, so these tests do not depend on @@ -205,7 +222,8 @@ describe('processImportQueue', () => { // no other assertion in this file fails if the write is deleted. const partWrites = dbWrite.huggingFaceImport.updateMany.mock.calls .map(([arg]) => arg) - .filter((arg: { data: Record }) => 'parts' in arg.data); + // Not the completion write, which also clears `parts` — these are the per-part heartbeats. + .filter((arg: { data: Record }) => 'parts' in arg.data && !arg.data.status); expect(partWrites).toHaveLength(3); for (const write of partWrites) { expect(write.data.heartbeatAt).toBeInstanceOf(Date); @@ -732,9 +750,9 @@ describe('unattached and delete', () => { }); it('scopes the lookup to the owner when the caller is not a moderator', async () => { - dbRead.huggingFaceImport.findFirst.mockResolvedValue(null); + dbWrite.huggingFaceImport.findFirst.mockResolvedValue(null); await expect(deleteImport({ id: 1, userId: 7, isModerator: false })).rejects.toThrow(); - expect(dbRead.huggingFaceImport.findFirst.mock.calls[0][0].where).toMatchObject({ + expect(dbWrite.huggingFaceImport.findFirst.mock.calls[0][0].where).toMatchObject({ id: 1, userId: 7, }); @@ -742,7 +760,7 @@ describe('unattached and delete', () => { it('refuses to delete an import that is still attached', async () => { // Deleting here would leave a model version pointing at bytes that no longer exist. - dbRead.huggingFaceImport.findFirst.mockResolvedValue({ + dbWrite.huggingFaceImport.findFirst.mockResolvedValue({ id: 1, status: 'Completed', bucket: 'b', @@ -760,7 +778,7 @@ describe('unattached and delete', () => { // The two-click data-loss path: detach leaves the ModelFile alive, so the row lands in the // unattached list while a published version is still serving those exact bytes. `modelFileId` // is a local pointer; the refcount over `ModelFile.url` is the global one. - dbRead.huggingFaceImport.findFirst.mockResolvedValue({ + dbWrite.huggingFaceImport.findFirst.mockResolvedValue({ id: 1, status: 'Completed', bucket: 'b2-transfer-bucket', @@ -779,7 +797,7 @@ describe('unattached and delete', () => { }); it('refuses to delete a transfer that is still running', async () => { - dbRead.huggingFaceImport.findFirst.mockResolvedValue({ + dbWrite.huggingFaceImport.findFirst.mockResolvedValue({ id: 1, status: 'Transferring', bucket: 'b', @@ -793,7 +811,7 @@ describe('unattached and delete', () => { }); it('frees the object before removing the row', async () => { - dbRead.huggingFaceImport.findFirst.mockResolvedValue({ + dbWrite.huggingFaceImport.findFirst.mockResolvedValue({ id: 1, status: 'Completed', // Deliberately NOT what the env resolves to: these bytes are in the bucket the transfer used. @@ -828,7 +846,7 @@ describe('unattached and delete', () => { it('aborts a live multipart before forgetting the row', async () => { // A Failed row can still hold an uploadId, and the row is the only handle that can free the // parts already uploaded — which are billed until something aborts them. - dbRead.huggingFaceImport.findFirst.mockResolvedValue({ + dbWrite.huggingFaceImport.findFirst.mockResolvedValue({ id: 1, status: 'Failed', bucket: 'b2-transfer-bucket', @@ -852,7 +870,7 @@ describe('unattached and delete', () => { }); it("aborts through a client that reaches the row's bucket", async () => { - dbRead.huggingFaceImport.findFirst.mockResolvedValue({ + dbWrite.huggingFaceImport.findFirst.mockResolvedValue({ id: 1, status: 'Failed', bucket: 'other-bucket', @@ -882,7 +900,7 @@ describe('unattached and delete', () => { }; it('keeps the row, and says why, when the multipart abort fails', async () => { - dbRead.huggingFaceImport.findFirst.mockResolvedValue(stuckRow); + dbWrite.huggingFaceImport.findFirst.mockResolvedValue(stuckRow); mockAbort.mockRejectedValue(new Error('The specified bucket does not exist')); const result = await deleteImport({ id: 1, userId: 7, isModerator: true }); @@ -899,7 +917,7 @@ describe('unattached and delete', () => { Object.assign(new Error(message), { name }); it('treats an upload that is already gone as removed, and needs no force', async () => { - dbRead.huggingFaceImport.findFirst.mockResolvedValue(stuckRow); + dbWrite.huggingFaceImport.findFirst.mockResolvedValue(stuckRow); mockAbort.mockRejectedValue(storageError('NoSuchUpload')); mockDeleteObject.mockRejectedValue(storageError('NoSuchKey')); dbWrite.huggingFaceImport.deleteMany.mockResolvedValue({ count: 1 }); @@ -911,7 +929,7 @@ describe('unattached and delete', () => { it('tries the other backend when the first does not know the bucket', async () => { // The row records only a bucket name; picking the wrong backend for it is how "The specified // bucket does not exist" happened. The upload is still there, on the other backend. - dbRead.huggingFaceImport.findFirst.mockResolvedValue({ ...stuckRow, bucket: 'other-bucket' }); + dbWrite.huggingFaceImport.findFirst.mockResolvedValue({ ...stuckRow, bucket: 'other-bucket' }); mockAbort .mockRejectedValueOnce(storageError('NoSuchBucket', 'The specified bucket does not exist')) .mockResolvedValueOnce(undefined); @@ -923,7 +941,7 @@ describe('unattached and delete', () => { }); it('asks for force only when no backend could remove it', async () => { - dbRead.huggingFaceImport.findFirst.mockResolvedValue({ ...stuckRow, bucket: 'other-bucket' }); + dbWrite.huggingFaceImport.findFirst.mockResolvedValue({ ...stuckRow, bucket: 'other-bucket' }); mockAbort.mockRejectedValue( storageError('NoSuchBucket', 'The specified bucket does not exist') ); @@ -936,7 +954,7 @@ describe('unattached and delete', () => { }); it('does not try another backend for an error that is not about the bucket', async () => { - dbRead.huggingFaceImport.findFirst.mockResolvedValue({ ...stuckRow, bucket: 'other-bucket' }); + dbWrite.huggingFaceImport.findFirst.mockResolvedValue({ ...stuckRow, bucket: 'other-bucket' }); mockAbort.mockRejectedValue(storageError('AccessDenied')); await deleteImport({ id: 1, userId: 7, isModerator: true }); @@ -945,7 +963,7 @@ describe('unattached and delete', () => { }); it('deletes a stuck row when forced, even though its upload cannot be aborted', async () => { - dbRead.huggingFaceImport.findFirst.mockResolvedValue(stuckRow); + dbWrite.huggingFaceImport.findFirst.mockResolvedValue(stuckRow); mockAbort.mockRejectedValue(new Error('The specified bucket does not exist')); dbWrite.huggingFaceImport.deleteMany.mockResolvedValue({ count: 1 }); @@ -959,7 +977,7 @@ describe('unattached and delete', () => { it('never forces past a model file that still points at the object', async () => { // Force overrides storage cleanup only — deleting live bytes is not a cleanup failure. - dbRead.huggingFaceImport.findFirst.mockResolvedValue({ + dbWrite.huggingFaceImport.findFirst.mockResolvedValue({ ...stuckRow, status: 'Completed', uploadId: null, @@ -976,7 +994,7 @@ describe('unattached and delete', () => { it('keeps the row when the object could not be deleted', async () => { // Otherwise the bytes stay in the bucket with nothing left pointing at them. - dbRead.huggingFaceImport.findFirst.mockResolvedValue({ + dbWrite.huggingFaceImport.findFirst.mockResolvedValue({ id: 1, status: 'Completed', bucket: 'b2-transfer-bucket', @@ -1014,7 +1032,7 @@ describe('retryImport', () => { const reset = () => dbWrite.huggingFaceImport.update.mock.calls[0]?.[0].data; it('restarts from nothing, in the backend configured now', async () => { - dbRead.huggingFaceImport.findFirst.mockResolvedValue(failedRow); + dbWrite.huggingFaceImport.findFirst.mockResolvedValue(failedRow); mockAbort.mockResolvedValue(undefined); expect(await retryImport({ id: 1, userId: 7, isModerator: true })).toEqual({ ok: true }); @@ -1032,7 +1050,7 @@ describe('retryImport', () => { it('asks before restarting when the old upload cannot be aborted', async () => { // Resuming the upload that cannot be aborted is what kept an import failing forever. - dbRead.huggingFaceImport.findFirst.mockResolvedValue(failedRow); + dbWrite.huggingFaceImport.findFirst.mockResolvedValue(failedRow); mockAbort.mockRejectedValue(new Error('The specified bucket does not exist')); const result = await retryImport({ id: 1, userId: 7, isModerator: true }); @@ -1042,7 +1060,7 @@ describe('retryImport', () => { }); it('restarts from nothing when forced past a failed abort', async () => { - dbRead.huggingFaceImport.findFirst.mockResolvedValue(failedRow); + dbWrite.huggingFaceImport.findFirst.mockResolvedValue(failedRow); mockAbort.mockRejectedValue(new Error('The specified bucket does not exist')); const result = await retryImport({ id: 1, userId: 7, isModerator: true, force: true }); @@ -1052,7 +1070,7 @@ describe('retryImport', () => { }); it('refuses a transfer that is still running', async () => { - dbRead.huggingFaceImport.findFirst.mockResolvedValue({ ...failedRow, status: 'Transferring' }); + dbWrite.huggingFaceImport.findFirst.mockResolvedValue({ ...failedRow, status: 'Transferring' }); const result = await retryImport({ id: 1, userId: 7, isModerator: true, force: true }); @@ -1126,3 +1144,428 @@ describe('renameGroup', () => { await expect(renameGroup(input)).rejects.toThrow(/No files found in group/); }); }); + +describe('auto-attach', () => { + const completedRow = (overrides: Record = {}) => ({ + attachVersionId: 42, + attachType: 'VAE', + userId: 7, + modelFileId: null, + status: 'Completed', + ...overrides, + }); + + /** What `buildAttachInput` reads once `attachIfRequested` decides to go ahead. */ + const attachableRow = () => ({ + id: 1, + filename: 'flux1-dev.safetensors', + url: 'https://s3.example/model-bucket/model/7/flux.safetensors', + key: 'model/7/flux.safetensors', + bucket: 'model-bucket', + sizeBytes: BigInt(2048), + status: 'Completed', + modelFileId: null, + }); + + beforeEach(() => { + vi.clearAllMocks(); + mockObjectExists.mockResolvedValue(true); + mockCreateModelFile.mockResolvedValue({ id: 900 }); + }); + + it('records where each file is headed, per file', async () => { + dbWrite.huggingFaceImport.createMany.mockResolvedValue({ count: 2 }); + dbWrite.huggingFaceImport.findMany.mockResolvedValue([ + { id: 1, filename: 'a.safetensors', status: 'Queued', modelFileId: null }, + { id: 2, filename: 'ae.safetensors', status: 'Queued', modelFileId: null }, + ]); + + const result = await enqueueImports({ + repo: { + repo: 'owner/name', + revision: 'abc123', + files: [ + { path: 'a.safetensors', size: 10, sha256: null }, + { path: 'ae.safetensors', size: 20, sha256: null }, + ], + } as never, + paths: ['a.safetensors', 'ae.safetensors'], + userId: 7, + attach: { modelVersionId: 42, types: { 'a.safetensors': 'Model', 'ae.safetensors': 'VAE' } }, + }); + + const queued = dbWrite.huggingFaceImport.createMany.mock.calls[0][0].data; + expect(queued).toMatchObject([ + { filename: 'a.safetensors', attachVersionId: 42, attachType: 'Model' }, + { filename: 'ae.safetensors', attachVersionId: 42, attachType: 'VAE' }, + ]); + // The QUERY, not the mock's answer: dropping `revision` returns another revision's rows, and + // the caller is handed import ids belonging to different bytes. + expect(dbWrite.huggingFaceImport.findMany).toHaveBeenCalledWith({ + where: { + repo: 'owner/name', + revision: 'abc123', + filename: { in: ['a.safetensors', 'ae.safetensors'] }, + }, + select: { + id: true, + filename: true, + status: true, + modelFileId: true, + attachVersionId: true, + }, + }); + expect(result.rows.map((row) => row.id)).toEqual([1, 2]); + }); + + it('leaves the attach columns null when nobody asked for one', async () => { + dbWrite.huggingFaceImport.createMany.mockResolvedValue({ count: 1 }); + dbWrite.huggingFaceImport.findMany.mockResolvedValue([]); + + await enqueueImports({ + repo: { + repo: 'owner/name', + revision: 'abc123', + files: [{ path: 'a.safetensors', size: 10, sha256: null }], + } as never, + paths: ['a.safetensors'], + userId: 7, + }); + + expect(dbWrite.huggingFaceImport.createMany.mock.calls[0][0].data[0]).toMatchObject({ + attachVersionId: null, + attachType: null, + }); + }); + + it('creates the model file and claims the import', async () => { + dbWrite.huggingFaceImport.findUnique.mockResolvedValue(completedRow()); + dbWrite.huggingFaceImport.findFirst.mockResolvedValue({ id: 1 }); + dbWrite.huggingFaceImport.findUniqueOrThrow.mockResolvedValue(attachableRow()); + dbWrite.huggingFaceImport.updateMany.mockResolvedValue({ count: 1 }); + + await attachIfRequested(1); + + expect(mockCreateModelFile).toHaveBeenCalledTimes(1); + const [call] = mockCreateModelFile.mock.calls; + expect(call[0].input).toMatchObject({ modelVersionId: 42, type: 'VAE', sizeKB: 2 }); + // 🔴 Moderator, deliberately: the job has no session, and every entry point that can set + // `attachVersionId` is moderator-gated. + expect(call[0]).toMatchObject({ userId: 7, isModerator: true }); + expect(dbWrite.huggingFaceImport.updateMany).toHaveBeenCalledWith({ + where: { id: 1, modelFileId: null }, + data: { modelFileId: 900, modelVersionId: 42 }, + }); + }); + + it('does nothing for an import nobody pointed at a version', async () => { + dbWrite.huggingFaceImport.findUnique.mockResolvedValue( + completedRow({ attachVersionId: null, attachType: null }) + ); + + await attachIfRequested(1); + + expect(mockCreateModelFile).not.toHaveBeenCalled(); + }); + + it('does nothing when the version is known but the type is not', async () => { + dbWrite.huggingFaceImport.findUnique.mockResolvedValue(completedRow({ attachType: null })); + + await attachIfRequested(1); + + expect(mockCreateModelFile).not.toHaveBeenCalled(); + }); + + it('does nothing for an import that is already attached', async () => { + dbWrite.huggingFaceImport.findUnique.mockResolvedValue(completedRow({ modelFileId: 900 })); + + await attachIfRequested(1); + + expect(mockCreateModelFile).not.toHaveBeenCalled(); + }); + + it('does nothing until the transfer has finished', async () => { + dbWrite.huggingFaceImport.findUnique.mockResolvedValue( + completedRow({ status: 'Transferring' }) + ); + + await attachIfRequested(1); + + expect(mockCreateModelFile).not.toHaveBeenCalled(); + }); + + it('records the reason and does not throw when the attach fails', async () => { + dbWrite.huggingFaceImport.findUnique.mockResolvedValue(completedRow()); + dbWrite.huggingFaceImport.findFirst.mockResolvedValue({ id: 1 }); + dbWrite.huggingFaceImport.findUniqueOrThrow.mockResolvedValue(attachableRow()); + mockCreateModelFile.mockRejectedValue(new Error('Model version not found')); + dbWrite.huggingFaceImport.updateMany.mockResolvedValue({ count: 1 }); + + await expect(attachIfRequested(1)).resolves.toBeUndefined(); + + expect(dbWrite.huggingFaceImport.updateMany).toHaveBeenCalledWith({ + where: { id: 1, modelFileId: null }, + data: { error: 'Attach failed: Model version not found' }, + }); + }); + + it('reports the created file when the import was claimed by someone else first', async () => { + dbWrite.huggingFaceImport.findUnique.mockResolvedValue(completedRow()); + dbWrite.huggingFaceImport.findFirst.mockResolvedValue({ id: 1 }); + dbWrite.huggingFaceImport.findUniqueOrThrow.mockResolvedValue(attachableRow()); + // Zero rows updated means the link lost the race; the file exists regardless, and its id is the + // only way to find it again. + dbWrite.huggingFaceImport.updateMany.mockResolvedValue({ count: 0 }); + + await attachIfRequested(1); + + const errorWrite = dbWrite.huggingFaceImport.updateMany.mock.calls + .map(([arg]) => arg) + .find((arg: { data: Record }) => typeof arg.data.error === 'string'); + expect(errorWrite?.data.error).toContain('900'); + }); +}); + +describe('the transfer attaches what it finished', () => { + /** The attach read, distinct from the pre-complete cancel re-read that precedes it. */ + const attachRow = { + attachVersionId: 42, + attachType: 'VAE', + userId: 7, + modelFileId: null, + status: 'Completed', + }; + + const attachable = { + id: 1, + filename: 'ae.safetensors', + url: 'https://s3.example/model-bucket/model/7/ae.safetensors', + key: 'model/7/ae.safetensors', + bucket: 'model-bucket', + sizeBytes: BigInt(2048), + status: 'Completed', + modelFileId: null, + }; + + function completingRow() { + // The cancel re-read first, then the attach read — one mock serves both, and a single + // `mockResolvedValue` would short-circuit the attach on `status !== 'Completed'`. + dbWrite.huggingFaceImport.findUnique + .mockResolvedValueOnce({ status: 'Transferring', claimedBy: 'test-worker' }) + .mockResolvedValue(attachRow); + dbWrite.huggingFaceImport.findFirst.mockResolvedValue({ id: 1 }); + dbWrite.huggingFaceImport.findUniqueOrThrow.mockResolvedValue(attachable); + } + + beforeEach(() => { + vi.clearAllMocks(); + mockCreateMultipart.mockResolvedValue('upload-1'); + mockUploadPart.mockImplementation( + async ({ partNumber }: { partNumber: number }) => `etag-${partNumber}` + ); + mockComplete.mockResolvedValue(undefined); + mockReadRange.mockImplementation( + async ({ start, end }: { start: number; end: number }) => new Uint8Array(end - start + 1) + ); + mockObjectExists.mockResolvedValue(true); + mockCreateModelFile.mockResolvedValue({ id: 900 }); + dbWrite.huggingFaceImport.update.mockResolvedValue({}); + dbWrite.huggingFaceImport.updateMany.mockResolvedValue({ count: 1 }); + dbWrite.huggingFaceImport.findMany.mockResolvedValue([]); + }); + + it('attaches the file the moment its last part lands', async () => { + // 🔴 The one line wiring auto-attach into the transfer. Everything else about the feature can be + // right and nothing will ever attach if this call goes. + claimOnce(baseRow({ sizeBytes: BigInt(PART_SIZE) })); + completingRow(); + + await processImportQueue({ + deadline: Date.now() + 5000, + worker: 'test-worker', + concurrency: 1, + partsInFlight: TEST_PARTS_IN_FLIGHT, + }); + + expect(mockCreateModelFile).toHaveBeenCalledTimes(1); + expect(mockCreateModelFile.mock.calls[0][0].input).toMatchObject({ + modelVersionId: 42, + type: 'VAE', + }); + }); + + it('does not attach when the completion write lost the row', async () => { + // The negative control for the test above: a cancel or another run took it, so the attach is + // theirs to do. + claimOnce(baseRow({ sizeBytes: BigInt(PART_SIZE) })); + completingRow(); + dbWrite.huggingFaceImport.updateMany.mockResolvedValue({ count: 0 }); + + await processImportQueue({ + deadline: Date.now() + 5000, + worker: 'test-worker', + concurrency: 1, + partsInFlight: TEST_PARTS_IN_FLIGHT, + }); + + expect(mockCreateModelFile).not.toHaveBeenCalled(); + }); + + it('sweeps up a file a dead run left unattached, under a claim', async () => { + dbWrite.huggingFaceImport.findMany.mockResolvedValue([{ id: 99 }]); + dbWrite.huggingFaceImport.findUnique.mockResolvedValue(attachRow); + dbWrite.huggingFaceImport.findFirst.mockResolvedValue({ id: 99 }); + dbWrite.huggingFaceImport.findUniqueOrThrow.mockResolvedValue(attachable); + dbWrite.$queryRaw.mockResolvedValue([]); + + await processImportQueue({ + deadline: Date.now() + 5000, + worker: 'test-worker', + concurrency: 1, + partsInFlight: TEST_PARTS_IN_FLIGHT, + }); + + expect(dbWrite.huggingFaceImport.findMany).toHaveBeenCalledWith({ + // A recorded failure is usually permanent, so the sweep only ever picks up crash gaps. + where: { + status: 'Completed', + modelFileId: null, + attachVersionId: { not: null }, + error: null, + }, + select: { id: true }, + orderBy: { completedAt: 'asc' }, + take: 1, + }); + // 🔴 Claimed first. Two runs passing the same read both mint a `ModelFile`, and the loser's is + // orphaned — `linkImportToFile` picks a winner only after both files exist. + expect(dbWrite.huggingFaceImport.updateMany).toHaveBeenCalledWith({ + where: { id: 99, modelFileId: null, claimedBy: null }, + data: { claimedBy: 'test-worker', claimedAt: expect.any(Date) }, + }); + expect(mockCreateModelFile).toHaveBeenCalledTimes(1); + }); + + it('files a tooling import under the system user', async () => { + dbWrite.huggingFaceImport.findUnique.mockResolvedValue({ ...attachRow, userId: null }); + dbWrite.huggingFaceImport.findFirst.mockResolvedValue({ id: 1 }); + dbWrite.huggingFaceImport.findUniqueOrThrow.mockResolvedValue(attachable); + dbWrite.huggingFaceImport.updateMany.mockResolvedValue({ count: 1 }); + + await attachIfRequested(1); + + expect(mockCreateModelFile.mock.calls[0][0].userId).toBe(IMPORT_SYSTEM_USER_ID); + // The value, not just the symbol: it is the `userId` segment of the stored object's key. + expect(IMPORT_SYSTEM_USER_ID).toBe(-1); + }); + + it('ends the auto-attach intent when a moderator detaches the file', async () => { + // 🔴 Otherwise the sweep matches the row on the next tick and re-attaches it — minting a SECOND + // `ModelFile`, since detach leaves the first one alive. + dbWrite.huggingFaceImport.findFirst.mockResolvedValue({ id: 1, modelFileId: 900 }); + dbWrite.huggingFaceImport.updateMany.mockResolvedValue({ count: 1 }); + + await detachImport({ id: 1, userId: 7, isModerator: true }); + + expect(dbWrite.huggingFaceImport.updateMany).toHaveBeenCalledWith({ + where: { id: 1, modelFileId: { not: null } }, + data: { + modelFileId: null, + modelVersionId: null, + attachVersionId: null, + attachType: null, + }, + }); + }); +}); + +describe('getImportStatus', () => { + beforeEach(() => vi.clearAllMocks()); + + it('emits sizes JSON can carry', async () => { + // `res.json()` throws on a BigInt, so dropping the conversion 500s every poll — the API's only + // progress surface. + dbWrite.huggingFaceImport.findUnique.mockResolvedValue({ + id: 1, + status: 'Transferring', + sizeBytes: BigInt(2048), + bytesTransferred: BigInt(1024), + }); + + const status = await getImportStatus(1); + + expect(JSON.stringify(status)).toContain('"sizeBytes":2048'); + expect(JSON.stringify(status)).toContain('"bytesTransferred":1024'); + }); + + it('reports nothing for an id that does not exist', async () => { + dbWrite.huggingFaceImport.findUnique.mockResolvedValue(null); + + await expect(getImportStatus(1)).resolves.toBeNull(); + }); +}); + +describe('reuseStoredFile', () => { + const input = { + repo: 'owner/name', + revision: 'abc123', + path: 'text_encoders/t5.safetensors', + sizeBytes: 4096, + sha256: 'DEADBEEF', + storedUrl: 'https://s3.example/model-bucket/model/9/t5.safetensors', + modelVersionId: 42, + type: 'Text Encoder' as const, + userId: 7, + }; + + beforeEach(() => { + vi.clearAllMocks(); + mockCreateModelFile.mockResolvedValue({ id: 900 }); + dbWrite.huggingFaceImport.create.mockResolvedValue({ id: 5 }); + dbWrite.huggingFaceImport.updateMany.mockResolvedValue({ count: 1 }); + }); + + it('attaches the bytes we already hold, without queueing a transfer', async () => { + dbWrite.huggingFaceImport.findUnique.mockResolvedValue(null); + + const result = await reuseStoredFile(input); + + expect(result).toEqual({ importId: 5, modelFileId: 900 }); + expect(mockCreateModelFile.mock.calls[0][0].input).toMatchObject({ + modelVersionId: 42, + type: 'Text Encoder', + name: 't5.safetensors', + url: input.storedUrl, + sizeKB: 4, + }); + // Queued for nothing: the row records the provenance and is Completed on arrival. + expect(dbWrite.huggingFaceImport.create.mock.calls[0][0].data).toMatchObject({ + status: 'Completed', + url: input.storedUrl, + attachVersionId: 42, + attachType: 'Text Encoder', + }); + }); + + it('claims no storage of its own', async () => { + // 🔴 These bytes are another file's object. A `key` here would let `deleteImport` reach for + // something this import never stored. + dbWrite.huggingFaceImport.findUnique.mockResolvedValue(null); + + await reuseStoredFile(input); + + const { data } = dbWrite.huggingFaceImport.create.mock.calls[0][0]; + expect(data.bucket).toBeUndefined(); + expect(data.key).toBeUndefined(); + }); + + it('does not import the same file twice', async () => { + dbWrite.huggingFaceImport.findUnique.mockResolvedValue({ id: 5, modelFileId: 900 }); + + const result = await reuseStoredFile(input); + + expect(result).toEqual({ importId: 5, modelFileId: 900 }); + expect(dbWrite.huggingFaceImport.create).not.toHaveBeenCalled(); + expect(mockCreateModelFile).not.toHaveBeenCalled(); + }); +}); diff --git a/src/server/services/huggingface-import.service.ts b/src/server/services/huggingface-import.service.ts index bd895f0127..09e2153b3d 100644 --- a/src/server/services/huggingface-import.service.ts +++ b/src/server/services/huggingface-import.service.ts @@ -1,5 +1,7 @@ import { Prisma } from '@prisma/client'; import { env } from '~/env/server'; +import { Tracker } from '~/server/clickhouse/client'; +import { createModelFile } from '~/server/controllers/model-file.controller'; import { dbRead, dbWrite } from '~/server/db/client'; import { logToAxiom } from '~/server/logging/client'; import { @@ -11,6 +13,7 @@ import { suggestFileType, type HuggingFaceRepo, } from '~/server/services/huggingface.service'; +import { constants } from '~/server/common/constants'; import { UploadType } from '~/server/common/enums'; import { abortMultipartUpload, @@ -43,6 +46,13 @@ const STALE_CLAIM_MINUTES = 20; * real database. (`minor-hash.service.ts` carries the same note for the same reason.) */ const STALE_CLAIM_INTERVAL = Prisma.raw(`make_interval(mins => ${STALE_CLAIM_MINUTES})`); +/** + * One per run. The sweep only ever has work after a pod died mid-attach, and an attach makes two + * external calls with no timeout of their own — so a backlog drains over a few minutes rather than + * spending the job lock in one tick. + */ +const SWEEP_SIZE = 1; + const MAX_ATTEMPTS = 5; const RETRY_BACKOFF_MINUTES = 5; /** @@ -94,6 +104,8 @@ const importSelect = { error: true, modelFileId: true, modelVersionId: true, + attachVersionId: true, + attachType: true, userId: true, createdAt: true, startedAt: true, @@ -262,14 +274,17 @@ export async function enqueueImports({ paths, userId, groupName, + attach, }: { repo: HuggingFaceRepo; paths: string[]; userId: number; groupName?: string; + /** Recorded per file so the transfer attaches itself — see `attachIfRequested`. */ + attach?: { modelVersionId: number; types: Record }; }) { const wanted = repo.files.filter((file) => paths.includes(file.path)); - if (!wanted.length) return { queued: 0, skipped: 0 }; + if (!wanted.length) return { queued: 0, skipped: 0, rows: [] }; // `(repo, revision, filename)` is unique, so re-queueing a repo adds only what is new. const result = await dbWrite.huggingFaceImport.createMany({ @@ -282,11 +297,201 @@ export async function enqueueImports({ sourceSha256: file.sha256, groupName: groupName?.trim() || defaultGroupName(repo.repo), userId, + attachVersionId: attach?.modelVersionId ?? null, + attachType: attach?.types[file.path] ?? null, })), skipDuplicates: true, }); - return { queued: result.count, skipped: wanted.length - result.count }; + // `createMany` returns a count, not rows, so the ids a caller polls need a second read. + const rows = await dbWrite.huggingFaceImport.findMany({ + where: { + repo: repo.repo, + revision: repo.revision, + filename: { in: wanted.map((file) => file.path) }, + }, + select: { id: true, filename: true, status: true, modelFileId: true, attachVersionId: true }, + }); + + return { queued: result.count, skipped: wanted.length - result.count, rows }; +} + +/** + * Puts a file we already store onto a version without transferring it again. + * + * Hugging Face publishes each LFS file's sha256 before any bytes move, so a byte-identical file we + * already hold — a text encoder shared by a dozen repos, say — is decided up front and costs nothing. + * + * 🔴 Keyed on the sha, never the filename. `ae.safetensors` and `model.safetensors` name different + * bytes in different repos, and attaching the wrong weights is invisible until someone generates. + * + * The row records the provenance and `bucket`/`key` stay null: these bytes are someone else's + * object, so deleting this import must never reach for them. + */ +export async function reuseStoredFile({ + repo, + revision, + path, + sizeBytes, + sha256, + storedUrl, + modelVersionId, + type, + userId, + groupName, +}: { + repo: string; + revision: string; + path: string; + sizeBytes: number | null; + sha256: string | null; + storedUrl: string; + modelVersionId: number; + type: ModelFileType; + userId: number; + groupName?: string; +}) { + const existingRow = await dbWrite.huggingFaceImport.findUnique({ + where: { repo_revision_filename: { repo, revision, filename: path } }, + select: { id: true, modelFileId: true }, + }); + if (existingRow) return { importId: existingRow.id, modelFileId: existingRow.modelFileId }; + + const row = await dbWrite.huggingFaceImport.create({ + data: { + repo, + revision, + filename: path, + sourceUrl: huggingFaceResolveUrl(repo, revision, path), + sizeBytes: sizeBytes === null ? null : BigInt(sizeBytes), + sourceSha256: sha256, + groupName: groupName?.trim() || defaultGroupName(repo), + userId, + status: 'Completed', + url: storedUrl, + completedAt: new Date(), + attachVersionId: modelVersionId, + attachType: type, + }, + select: { id: true }, + }); + + const name = path.split('/').pop() ?? path; + const file = await createModelFile({ + input: { + modelVersionId, + type, + name, + url: storedUrl, + sizeKB: bytesToKB(sizeBytes ?? 0), + // No `backend`/`s3Path`: the controller derives both from the url, which is what registers + // this new file id against the object the original upload put there. + metadata: { format: getModelFileFormat(name) }, + }, + userId, + isModerator: true, + track: new Tracker(), + }); + await linkImportToFile({ id: row.id, modelFileId: file.id, modelVersionId }); + + return { importId: row.id, modelFileId: file.id }; +} + +/** Imports asked for by tooling rather than a person. Not a real `User` row — `userId` has no FK. */ +export const IMPORT_SYSTEM_USER_ID = constants.system.user.id; + +/** Not owner-scoped, unlike `ownedImport` — only reachable behind `WEBHOOK_TOKEN`. */ +export async function getImportStatus(id: number) { + // Primary: a caller polls the id the POST just handed it, which the replica may not have yet. + const row = await dbWrite.huggingFaceImport.findUnique({ + where: { id }, + select: { + id: true, + repo: true, + revision: true, + filename: true, + status: true, + sizeBytes: true, + bytesTransferred: true, + attachVersionId: true, + modelVersionId: true, + modelFileId: true, + error: true, + completedAt: true, + }, + }); + if (!row) return null; + // BigInt does not survive JSON, and these are file sizes: a number is exact well past any of them. + return { + ...row, + sizeBytes: row.sizeBytes === null ? null : Number(row.sizeBytes), + bytesTransferred: Number(row.bytesTransferred), + }; +} + +/** + * 🔴 Runs as a moderator — the job has no session, so the version is NOT re-checked against the + * requester. Any surface that lets a non-moderator set `attachVersionId` must check that ownership + * at enqueue time. + * + * Never throws: the bytes are already transferred, and an unattached row with its reason recorded is + * exactly what the Unattached tab surfaces. + */ +export async function attachIfRequested(id: number) { + const row = await dbWrite.huggingFaceImport.findUnique({ + where: { id }, + select: { + attachVersionId: true, + attachType: true, + userId: true, + modelFileId: true, + status: true, + }, + }); + if (!row || row.modelFileId || row.status !== 'Completed') return; + if (!row.attachVersionId || !row.attachType) return; + + const userId = row.userId ?? IMPORT_SYSTEM_USER_ID; + try { + const { importId, ...fileInput } = await buildAttachInput({ + id, + modelVersionId: row.attachVersionId, + type: row.attachType as ModelFileType, + userId, + isModerator: true, + }); + const file = await createModelFile({ + input: fileInput, + userId, + isModerator: true, + track: new Tracker(), + }); + const linked = await linkImportToFile({ + id: importId, + modelFileId: file.id, + modelVersionId: row.attachVersionId, + }); + if (!linked) + throw new Error( + `Created model file ${file.id}, but this import was attached by someone else first.` + ); + } catch (error) { + const message = error instanceof Error ? error.message : String(error); + logToAxiom({ + type: 'error', + name: 'huggingface-import', + message: 'auto-attach failed; the file is transferred but unattached', + importId: id, + modelVersionId: row.attachVersionId, + error: message, + }); + await dbWrite.huggingFaceImport + .updateMany({ + where: { id, modelFileId: null }, + data: { error: `Attach failed: ${message}` }, + }) + .catch(() => undefined); + } } /** Scoped by owner as well as id — the page is moderator-only today, the service is not. */ @@ -299,7 +504,10 @@ async function ownedImport({ userId: number; isModerator: boolean; }) { - const row = await dbRead.huggingFaceImport.findFirst({ + // Primary: the transfer job attaches microseconds after its own completion write, and a replica + // that has not caught up reports the row as still transferring — which the caller records as a + // permanent failure. + const row = await dbWrite.huggingFaceImport.findFirst({ where: { id, userId: isModerator ? undefined : userId }, select: { id: true, @@ -337,7 +545,7 @@ export async function buildAttachInput({ // predicate for the sake of a wider `select`, which left two copies of the rule — and this is the // copy that mints a `ModelFile` on a caller-supplied version, so it is the worst one to let drift. await ownedImport({ id, userId, isModerator }); - const row = await dbRead.huggingFaceImport.findUniqueOrThrow({ + const row = await dbWrite.huggingFaceImport.findUniqueOrThrow({ where: { id }, select: { id: true, @@ -418,7 +626,10 @@ export async function detachImport(input: { id: number; userId: number; isModera const row = await ownedImport(input); const { count } = await dbWrite.huggingFaceImport.updateMany({ where: { id: row.id, modelFileId: { not: null } }, - data: { modelFileId: null, modelVersionId: null }, + // 🔴 The attach target goes too. Detaching is a decision that this file does not belong on that + // version, and `retryPendingAttachments` would otherwise re-attach it on the next tick — minting + // a SECOND `ModelFile`, since the first one detach leaves alive. + data: { modelFileId: null, modelVersionId: null, attachVersionId: null, attachType: null }, }); return { ok: count === 1 }; } @@ -774,7 +985,9 @@ async function advanceImport( if (!uploadId || !key) { // Refuse rather than invent an owner: `model/0/…` is a key no upload path could produce, and the - // userId segment is what `/api/upload/sign-part` authorises against. + // userId segment is what `/api/upload/sign-part` authorises against. `IMPORT_SYSTEM_USER_ID` is + // the one sanctioned exception — a server transfer never presents a part to `sign-part`. `IMPORT_SYSTEM_USER_ID` is + // the one sanctioned exception — a server transfer never presents a part to `sign-part`. if (!row.userId) throw new Error(`Import ${row.id} has no owner; refusing to build a key for it`); key = buildUploadKey( @@ -878,7 +1091,7 @@ async function advanceImport( // Deliberately not `getCustomPutUrl` — that bumps `recordB2PresignIssued`, a counter whose whole // purpose is measuring browser-direct uploads, and a server transfer is not one. const { url } = await getGetUrlByKey(key, { s3, bucket }); - await dbWrite.huggingFaceImport.updateMany({ + const { count: completed } = await dbWrite.huggingFaceImport.updateMany({ where: { id: row.id, claimedBy: row.claimedBy, status: { not: 'Canceled' } }, data: { status: 'Completed', @@ -890,15 +1103,56 @@ async function advanceImport( // Spent. A retained id makes every later abort fail against a finished upload, burying the // one abort failure that means parts are still billed. uploadId: null, + // The resume ledger, only read while a transfer is in flight. Measured at ~40 bytes per 16 MB + // part, so a finished 50 GB import would otherwise carry ~125 KB of dead JSON on a row that + // every later query has to read past. + parts: Prisma.DbNull, }, }); + + // count 0 means a cancel or another run took the row — its attach is theirs to do. + if (completed) await attachIfRequested(row.id); return movedBytes; } +/** + * Files that finished but never got attached, because the pod died between the two writes. + * + * A null error keeps this a crash sweep: a recorded failure is usually permanent — a deleted version, + * a type the extension refuses — and retrying those would fill every slot forever. + */ +async function retryPendingAttachments(worker: string, deadline: number) { + const pending = await dbWrite.huggingFaceImport.findMany({ + where: { status: 'Completed', modelFileId: null, attachVersionId: { not: null }, error: null }, + select: { id: true }, + orderBy: { completedAt: 'asc' }, + take: SWEEP_SIZE, + }); + + for (const row of pending) { + if (Date.now() >= deadline) return; + // 🔴 The same claim the transfer takes. An attach is a read, then an S3 head, then two awaited + // external calls, then a write — long enough for a second run to pass the same read and mint a + // second `ModelFile` that nothing would ever delete. `linkImportToFile` picks a winner, but only + // after both files exist. + const { count } = await dbWrite.huggingFaceImport.updateMany({ + where: { id: row.id, modelFileId: null, claimedBy: null }, + data: { claimedBy: worker, claimedAt: new Date() }, + }); + if (!count) continue; + try { + await attachIfRequested(row.id); + } finally { + await yieldClaim(row.id).catch(() => undefined); + } + } +} + /** * Drains the queue until `deadline`. Bounded work per call by design: a transfer is a sequence of * resumable parts, so the job never needs a run longer than its own lock. */ + export async function processImportQueue({ deadline, worker, @@ -928,6 +1182,10 @@ export async function processImportQueue({ } }; + // Before the drain, not after: the drain runs until the deadline by construction, so anything + // queued behind it spends time the job's lock does not have. + await retryPendingAttachments(worker, deadline); + const startedAt = Date.now(); await Promise.all(Array.from({ length: concurrency }, drain)); const seconds = (Date.now() - startedAt) / 1000; diff --git a/src/server/services/orchestrator/__tests__/createModelFileScanRequest.test.ts b/src/server/services/orchestrator/__tests__/createModelFileScanRequest.test.ts index a79f615dcb..cd2b3b8cfb 100644 --- a/src/server/services/orchestrator/__tests__/createModelFileScanRequest.test.ts +++ b/src/server/services/orchestrator/__tests__/createModelFileScanRequest.test.ts @@ -395,6 +395,23 @@ describe('createModelFileScanRequest', () => { expect(mockSubmitWorkflow).toHaveBeenCalled(); }); + it('gives the submit a deadline of its own', async () => { + // No `wait` is passed, so this is an enqueue: without a signal it inherits undici's 300s + // default while the upload response and the import job's lock both wait on it. + mockResolveDownloadUrl.mockResolvedValueOnce({ url: 'https://cdn/x' }); + mockSubmitWorkflow.mockResolvedValue({ + data: { id: 'wf-1' }, + error: undefined, + response: { status: 200 }, + }); + + await createModelFileScanRequest(baseInput); + + const [submitted] = mockSubmitWorkflow.mock.calls[0]; + expect(submitted.signal).toBeInstanceOf(AbortSignal); + expect(submitted.query?.wait).toBeUndefined(); + }); + it('retries pre-flight once after a 60s wait when the first attempt fails (sync-lag tolerance)', async () => { mockResolveDownloadUrl .mockRejectedValueOnce(new Error('not in resolver yet')) diff --git a/src/server/services/orchestrator/orchestrator.service.ts b/src/server/services/orchestrator/orchestrator.service.ts index ab9568f2f8..b5f1b74d10 100644 --- a/src/server/services/orchestrator/orchestrator.service.ts +++ b/src/server/services/orchestrator/orchestrator.service.ts @@ -44,6 +44,17 @@ import { // explicitly wants to block for the workflow, so we must not abort it early). const IMAGE_INGEST_SUBMIT_ATTEMPT_TIMEOUT_MS = 15_000; +/** + * The scan submit passes no `wait`, so this is an enqueue: the orchestrator accepts the workflow and + * returns. Without a signal it inherits undici's 300s default, and every caller awaits it — the + * upload response, and the import job inside its own lock. Sized like its sibling above rather than + * to a measured P99, which nothing records for this call. + * + * A fired timeout throws, which every caller already treats as a transient failure: `scanRequestedAt` + * stays null and `scanFilesFallbackJob` re-submits within five minutes. + */ +const MODEL_FILE_SCAN_SUBMIT_TIMEOUT_MS = 15_000; + const IMAGE_TAGGING_MODEL = 'urn:air:siglip2:repository:huggingface:cella110n/cl_tagger_v2@b57909b8e9c63f71e208a26473e7aabdf45ed6b6.tar'; const IMAGE_TAGGING_THRESHOLD = 0.55; @@ -781,6 +792,7 @@ export async function createModelFileScanRequest({ const { data, error, response } = await submitWorkflow({ client: internalOrchestratorClient, + signal: AbortSignal.timeout(MODEL_FILE_SCAN_SUBMIT_TIMEOUT_MS), body: { metadata, currencies: [], diff --git a/src/utils/__tests__/storage-resolver.timeout.test.ts b/src/utils/__tests__/storage-resolver.timeout.test.ts new file mode 100644 index 0000000000..21cb527623 --- /dev/null +++ b/src/utils/__tests__/storage-resolver.timeout.test.ts @@ -0,0 +1,53 @@ +import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest'; +import '~/__tests__/mocks/logging.mock'; +import { setEnv } from '~/__tests__/mocks/env.mock'; +import { registerFileLocation } from '~/utils/storage-resolver'; + +/** + * Without a signal this inherits undici's 300s default. Every caller awaits it, so one hung + * registration holds the caller's job lock far past its own budget — which is what lets a second + * run start on work the first still owns. + */ +const params = { + fileId: 1, + modelVersionId: 42, + modelId: 5, + backend: 'backblaze', + path: 'model/7/x.safetensors', + sizeKb: 2, +}; + +const fetchMock = vi.fn(); + +beforeEach(() => { + vi.stubGlobal('fetch', fetchMock); + fetchMock.mockReset().mockResolvedValue({ ok: true, json: async () => ({}) }); + setEnv({ + STORAGE_RESOLVER_INTERNAL_URL: 'https://resolver.example', + STORAGE_RESOLVER_INTERNAL_TOKEN: 'token', + }); +}); + +afterEach(() => vi.unstubAllGlobals()); + +describe('registerFileLocation', () => { + it('gives up rather than hanging its caller', async () => { + await registerFileLocation(params); + + const [, init] = fetchMock.mock.calls[0]; + expect(init.signal).toBeInstanceOf(AbortSignal); + // Generous for an internal write, and far inside any caller's lock. + expect(init.signal.aborted).toBe(false); + }); + + it('does not call out at all when the resolver is not configured', async () => { + setEnv({ + STORAGE_RESOLVER_INTERNAL_URL: undefined, + STORAGE_RESOLVER_INTERNAL_TOKEN: undefined, + }); + + await registerFileLocation(params); + + expect(fetchMock).not.toHaveBeenCalled(); + }); +}); diff --git a/src/utils/storage-resolver.ts b/src/utils/storage-resolver.ts index 4f5e5bb1f4..fbb95c2eeb 100644 --- a/src/utils/storage-resolver.ts +++ b/src/utils/storage-resolver.ts @@ -1,6 +1,12 @@ import { env } from '~/env/server'; import { logToAxiom } from '~/server/logging/client'; +/** + * Shorter than the 30s the deregister calls use: this is one small write on the upload path, which a + * caller awaits before responding, while those are bulk post-commit cleanups nobody waits on. + */ +const REGISTER_TIMEOUT_MS = 10_000; + export async function registerFileLocation(params: { fileId: number; modelVersionId: number; @@ -28,6 +34,7 @@ export async function registerFileLocation(params: { Authorization: `Bearer ${env.STORAGE_RESOLVER_INTERNAL_TOKEN}`, }, body: JSON.stringify(params), + signal: AbortSignal.timeout(REGISTER_TIMEOUT_MS), }); if (!response.ok) {