From 9f63c1ee02079394d0d781f945861cfc20318935 Mon Sep 17 00:00:00 2001 From: andreip135 <338078672+andreip135@users.noreply.github.com> Date: Thu, 8 Oct 2026 14:14:30 +0300 Subject: [PATCH 1/8] check Comfy model download status --- src/@types/C2D/ServiceOnDemand.ts | 1 + src/components/c2d/comfyDownload.ts | 179 ++++++++++++++++++++ src/components/c2d/compute_engine_docker.ts | 16 +- src/components/c2d/modelDownload.ts | 23 ++- src/components/c2d/serviceEngines.ts | 10 ++ src/test/unit/service/comfyDownload.test.ts | 152 +++++++++++++++++ 6 files changed, 374 insertions(+), 7 deletions(-) create mode 100644 src/components/c2d/comfyDownload.ts create mode 100644 src/test/unit/service/comfyDownload.test.ts diff --git a/src/@types/C2D/ServiceOnDemand.ts b/src/@types/C2D/ServiceOnDemand.ts index aa3b4a681..6ee02c850 100644 --- a/src/@types/C2D/ServiceOnDemand.ts +++ b/src/@types/C2D/ServiceOnDemand.ts @@ -193,6 +193,7 @@ export interface ServiceModelDownload { percent?: number filesComplete: number filesInFlight: number // the hub fetches files in parallel, so only the aggregate is meaningful + filesTotal?: number // files the service will download, when it lists them up front (ComfyUI) updatedAt: number // Unix ms } diff --git a/src/components/c2d/comfyDownload.ts b/src/components/c2d/comfyDownload.ts new file mode 100644 index 000000000..04b760a49 --- /dev/null +++ b/src/components/c2d/comfyDownload.ts @@ -0,0 +1,179 @@ +import type Dockerode from 'dockerode' +import { pipeline } from 'stream' +import * as tarStream from 'tar-stream' +import type { ServiceModelDownload } from '../../@types/C2D/ServiceOnDemand.js' +import { + buildModelDownload, + cachedHubLookup, + fileSize, + HF_TIMEOUT_MS +} from './modelDownload.js' + +/** + * Model-download progress for the ComfyUI bundle templates. + * + * Before fetching anything, a bundle's script lists the weights its workflows need in + * `/.models.tsv`, one `\t` line each. It then downloads them with curl to + * `/models//.part`, renaming each to `` once its size checks out. `` is + * `/data/outputs/comfy` when a bucket is mounted, `/tmp/comfy` otherwise. + * + * The script is client-supplied, so the list is untrusted: only plain Hugging Face file URLs are + * kept, and those are the only addresses the node ever contacts on its behalf. + */ + +const MODEL_LIST_FILE = '.models.tsv' +const MAX_LIST_BYTES = 64 * 1024 +const MAX_ENTRIES = 64 +const MODEL_DIR = /^[a-z0-9_]{1,32}$/ +const URL_SEGMENT = /^[\w.-]+$/ + +export interface ComfyModelEntry { + url: string + path: string +} + +export function comfyBaseDir(hasBucket: boolean): string { + return hasBucket ? '/data/outputs/comfy' : '/tmp/comfy' +} + +/** `https://huggingface.co///resolve//.safetensors`, exactly. */ +function isHubFileUrl(value: string): boolean { + let url: URL + try { + url = new URL(value) + } catch { + return false + } + // Rebuilt from the parsed path, so a query, credentials, a port or a normalized `..` all fail. + if (value !== `https://huggingface.co${url.pathname}`) { + return false + } + const segments = url.pathname.split('/').slice(1) + return ( + segments.length >= 5 && + segments[2] === 'resolve' && + segments.every((segment) => URL_SEGMENT.test(segment)) && + url.pathname.endsWith('.safetensors') + ) +} + +export function parseComfyModelList(text: string, base: string): ComfyModelEntry[] { + const entries = new Map() + for (const line of text.split('\n')) { + if (entries.size >= MAX_ENTRIES) { + break + } + const [dir, url] = line.trim().split('\t') + if (!MODEL_DIR.test(dir) || !url || !isHubFileUrl(url)) { + continue + } + // Same destination the script derives: `$MODELS/$sub/$(basename "$url")`. + const path = `${base}/models/${dir}/${url.slice(url.lastIndexOf('/') + 1)}` + if (!entries.has(path)) { + entries.set(path, { url, path }) + } + } + return [...entries.values()] +} + +/** One small file from the container, via Docker's archive endpoint. Null when absent or too big. */ +async function readContainerFile( + container: Dockerode.Container, + path: string +): Promise<{ text: string; mtime: number } | null> { + try { + const extract = tarStream.extract() + pipeline(await container.getArchive({ path }), extract, () => {}) + for await (const entry of extract) { + const { type, size, mtime } = entry.header + if (type !== 'file' || (size ?? 0) > MAX_LIST_BYTES) { + return null + } + const chunks: Buffer[] = [] + for await (const chunk of entry) { + chunks.push(chunk) + } + return { + text: Buffer.concat(chunks).toString('utf8'), + mtime: mtime?.getTime() ?? 0 + } + } + return null + } catch { + return null + } +} + +/** + * A file's size from the Hub without downloading it. A file in large-file storage answers with a + * redirect to the CDN whose `X-Linked-Size` is the file's size; a small one is served directly. + */ +async function lookupHubFileBytes(url: string): Promise { + try { + const response = await fetch(url, { + method: 'HEAD', + redirect: 'manual', + signal: AbortSignal.timeout(HF_TIMEOUT_MS) + }) + if (response.status >= 400) { + // Missing, or gated without a token: asking again will not change the answer. + return [401, 403, 404].includes(response.status) ? null : undefined + } + const size = Number( + response.headers.get('x-linked-size') ?? + (response.ok ? response.headers.get('content-length') : null) + ) + return size > 0 ? size : null + } catch { + return undefined + } +} + +/** + * Bytes on disk against the size of everything the list names. Null until the list exists — and, + * in a reused bucket, until this launch has rewritten the previous one (hence the mtime check). + * Never throws. + */ +export async function sampleComfyModelDownload( + container: Dockerode.Container, + hasBucket: boolean, + startedAt: number +): Promise { + const base = comfyBaseDir(hasBucket) + const list = await readContainerFile(container, `${base}/${MODEL_LIST_FILE}`) + // tar keeps whole seconds, hence the one-second slack. + if (!list || list.mtime < startedAt - 1000) { + return null + } + const entries = parseComfyModelList(list.text, base) + if (entries.length === 0) { + return null + } + const sizes = await Promise.all( + entries.map(({ url }) => cachedHubLookup(url, () => lookupHubFileBytes(url))) + ) + let downloadedBytes = 0 + let files = 0 + let inFlight = 0 + for (const { path } of entries) { + // Partial first: a file renamed between the two checks is then still found finished. + const partial = await fileSize(container, `${path}.part`) + if (partial !== null) { + downloadedBytes += partial + inFlight++ + continue + } + const finished = await fileSize(container, path) + if (finished !== null) { + downloadedBytes += finished + files++ + } + } + const totalBytes = sizes.every((size) => size !== null) + ? sizes.reduce((sum, size) => sum + size, 0) + : null + return { + ...buildModelDownload({ downloadedBytes, files, inFlight }, totalBytes, null), + filesTotal: entries.length + } +} diff --git a/src/components/c2d/compute_engine_docker.ts b/src/components/c2d/compute_engine_docker.ts index bec9f76a3..d4b32171b 100755 --- a/src/components/c2d/compute_engine_docker.ts +++ b/src/components/c2d/compute_engine_docker.ts @@ -99,6 +99,7 @@ import { isModelDownloadComplete, ModelDownloadSampler } from './modelDownload.js' +import { sampleComfyModelDownload } from './comfyDownload.js' import type { DockerMountObject } from '../../@types/PersistentStorage.js' import { resolveServiceImage } from './serviceResourceMatching.js' import { @@ -4817,7 +4818,7 @@ export class C2DEngineDocker extends C2DEngine { // when it says nothing about why. const modelDownload = everReady ? undefined - : await this.sampleModelDownload(job, engine) + : await this.sampleModelDownload(job, engine, startedAt) // Same cross-process guard as the metrics write: a lifecycle transition must win. if ( @@ -4857,9 +4858,10 @@ export class C2DEngineDocker extends C2DEngine { */ private async sampleModelDownload( job: ServiceJob, - engine: ServiceEngineProfile + engine: ServiceEngineProfile, + startedAt: number ): Promise { - if (!engine.modelCachePath) return undefined + if (!engine.modelCachePath && !engine.comfyModelList) return undefined // Downloaded: what remains is the engine loading weights, which the cache says nothing about. // Returning nothing keeps the stored 100% record as it is. if (isModelDownloadComplete(job.modelDownload)) { @@ -4867,6 +4869,14 @@ export class C2DEngineDocker extends C2DEngine { } try { const container = this.docker.getContainer(job.containerId) + if (engine.comfyModelList) { + const download = await sampleComfyModelDownload( + container, + !!job.outputBucketId, + startedAt + ) + return download ?? undefined + } const downloaded = await this.modelDownloadSampler.sample(job.serviceId, container) if (!downloaded) return undefined // Only a Hugging Face repo has a size the node can look up; a local path or an object-store diff --git a/src/components/c2d/modelDownload.ts b/src/components/c2d/modelDownload.ts index 073f1cc23..e368a84fe 100644 --- a/src/components/c2d/modelDownload.ts +++ b/src/components/c2d/modelDownload.ts @@ -73,7 +73,7 @@ async function statContainerPath( * `huggingface_hub` can link a blob into a shared store — the link's own size would count a few * bytes for a multi-gigabyte file. */ -async function fileSize( +export async function fileSize( container: Dockerode.Container, path: string ): Promise { @@ -328,7 +328,7 @@ const BYTES_PER_PARAM: Record = { } const HF_MODEL_API = 'https://huggingface.co/api/models' -const HF_TIMEOUT_MS = 8000 +export const HF_TIMEOUT_MS = 8000 // One lookup per model for the life of the process: the answer cannot change for a given repo, and // this is read on the metrics cadence for every starting service. const totalBytesCache = new Map() @@ -355,12 +355,23 @@ export async function fetchModelTotalBytes( quant?: string ): Promise { const key = quant ? `${modelId}:${quant}` : modelId + return await cachedHubLookup(key, () => lookupModelTotalBytes(modelId, quant)) +} + +/** + * Runs a Hub size lookup at most once per key: a size or a definitive null is kept for the life of + * the process, while undefined (the Hub could not be asked) is retried after HUB_RETRY_AFTER_MS. + */ +export async function cachedHubLookup( + key: string, + lookup: () => Promise +): Promise { if (totalBytesCache.has(key)) return totalBytesCache.get(key) ?? null const failedAt = hubFailedAt.get(key) if (failedAt !== undefined && Date.now() - failedAt < HUB_RETRY_AFTER_MS) { return null } - const total = await lookupModelTotalBytes(modelId, quant) + const total = await lookup() if (total === undefined) { hubFailedAt.set(key, Date.now()) return null @@ -470,11 +481,15 @@ async function fetchGgufFileBytes( * Whether a recorded download has finished: the total reached AND nothing still arriving. The total * is an estimate (dtype × parameters), so on its own it can read 100% while the last shard is still * being written; no partial files left is what makes it final. A record with no total never - * completes here, and sampling simply carries on until the service is ready. + * completes here, and sampling simply carries on until the service is ready. A record that knows + * its file count (a ComfyUI model list) is complete once every listed file is. */ export function isModelDownloadComplete( download: ServiceModelDownload | undefined ): boolean { + if (download?.filesTotal !== undefined) { + return download.filesComplete >= download.filesTotal + } return download?.percent === 100 && download.filesInFlight === 0 } diff --git a/src/components/c2d/serviceEngines.ts b/src/components/c2d/serviceEngines.ts index 621e6e000..3f5a70c75 100644 --- a/src/components/c2d/serviceEngines.ts +++ b/src/components/c2d/serviceEngines.ts @@ -45,6 +45,8 @@ export interface ServiceEngineProfile { * for a GGUF repo carrying a dozen quantizations differs by an order of magnitude. */ modelQuantFromCommand?: (cmd: string[] | undefined) => string | null + /** Progress is read from the model list a ComfyUI bundle script writes (see comfyDownload). */ + comfyModelList?: boolean } /** @@ -149,6 +151,14 @@ export const SERVICE_ENGINE_PROFILES: ServiceEngineProfile[] = [ modelCachePath: '/root/.cache/huggingface/hub', modelIdFromCommand: hfRepoIdFromLlamaCppCommand, modelQuantFromCommand: quantFromLlamaCppCommand + }, + { + id: 'comfyui', + matchesImage: (image) => /(^|\/)yanwk\/comfyui-boot$/.test(image), + // ComfyUI binds its port only after the bundle script has fetched its models and every custom + // node has loaded, so the first 200 means the UI is usable. + probe: { path: '/system_stats', port: 8188, expectStatus: [200] }, + comfyModelList: true } ] diff --git a/src/test/unit/service/comfyDownload.test.ts b/src/test/unit/service/comfyDownload.test.ts new file mode 100644 index 000000000..da7abc5f8 --- /dev/null +++ b/src/test/unit/service/comfyDownload.test.ts @@ -0,0 +1,152 @@ +import { expect } from 'chai' +import * as tarStream from 'tar-stream' +import type { ServiceJob } from '../../../@types/C2D/ServiceOnDemand.js' +import { resolveServiceEngine } from '../../../components/c2d/serviceEngines.js' +import { isModelDownloadComplete } from '../../../components/c2d/modelDownload.js' +import { + parseComfyModelList, + sampleComfyModelDownload +} from '../../../components/c2d/comfyDownload.js' + +const HF = 'https://huggingface.co' + +describe('ComfyUI engine profile', () => { + const job = (image: string) => ({ image }) as ServiceJob + + it('recognizes the comfyui-boot image under any tag', () => { + expect(resolveServiceEngine(job('yanwk/comfyui-boot'))?.id).to.equal('comfyui') + expect( + resolveServiceEngine(job('yanwk/comfyui-boot:cu130-megapak-pt211'))?.id + ).to.equal('comfyui') + expect(resolveServiceEngine(job('someone/comfyui-boot'))).to.equal(null) + }) +}) + +describe('parseComfyModelList', () => { + it('keeps plain Hugging Face file URLs and derives the script destination', () => { + const text = [ + `vae\t${HF}/org/repo/resolve/main/split_files/vae/a.safetensors`, + `loras\t${HF}/org/repo/resolve/main/b.safetensors`, + `loras\t${HF}/org/repo/resolve/main/b.safetensors`, + '' + ].join('\n') + expect(parseComfyModelList(text, '/tmp/comfy')).to.deep.equal([ + { + url: `${HF}/org/repo/resolve/main/split_files/vae/a.safetensors`, + path: '/tmp/comfy/models/vae/a.safetensors' + }, + { + url: `${HF}/org/repo/resolve/main/b.safetensors`, + path: '/tmp/comfy/models/loras/b.safetensors' + } + ]) + }) + + it('drops anything that is not a plain Hugging Face file URL or a safe directory', () => { + const text = [ + `../etc\t${HF}/org/repo/resolve/main/a.safetensors`, + `vae\thttps://example.com/org/repo/resolve/main/a.safetensors`, + `vae\t${HF}/org/repo/resolve/main/a.safetensors?download=true`, + `vae\t${HF}:8443/org/repo/resolve/main/a.safetensors`, + `vae\t${HF}/org/repo/resolve/main/../a.safetensors`, + `vae\t${HF}/org/repo/blob/main/a.safetensors`, + `vae\t${HF}/org/repo/resolve/main/a.bin`, + 'vae' + ].join('\n') + expect(parseComfyModelList(text, '/tmp/comfy')).to.deep.equal([]) + }) +}) + +describe('sampleComfyModelDownload', () => { + const realFetch = globalThis.fetch + const startedAt = Date.parse('2026-10-07T10:00:00Z') + + afterEach(() => { + globalThis.fetch = realFetch + }) + + // A container double: `getArchive()` returns the list as a tar, `infoArchive()` answers stats. + function fakeContainer(list: string, mtime: Date, files: Record) { + return { + getArchive: () => { + const pack = tarStream.pack() + pack.entry({ name: '.models.tsv', mtime }, list) + pack.finalize() + return Promise.resolve(pack) + }, + infoArchive: ({ path }: { path: string }) => { + if (files[path] === undefined) { + return Promise.reject(new Error('not found')) + } + const header = Buffer.from(JSON.stringify({ size: files[path] })).toString( + 'base64' + ) + return Promise.resolve({ headers: { 'x-docker-container-path-stat': header } }) + } + } as any + } + + function hubSizes(sizes: Record) { + globalThis.fetch = ((url: string) => + Promise.resolve( + new Response(null, { + status: 302, + headers: { 'x-linked-size': String(sizes[url]) } + }) + )) as typeof fetch + } + + it('measures finished and partial files against the listed total', async () => { + const a = `${HF}/sample-test/repo/resolve/main/a.safetensors` + const b = `${HF}/sample-test/repo/resolve/main/b.safetensors` + const c = `${HF}/sample-test/repo/resolve/main/c.safetensors` + hubSizes({ [a]: 1000, [b]: 2000, [c]: 1000 }) + const container = fakeContainer( + `vae\t${a}\nloras\t${b}\nloras\t${c}\n`, + new Date(startedAt + 5000), + { + '/data/outputs/comfy/models/vae/a.safetensors': 1000, + '/data/outputs/comfy/models/loras/b.safetensors.part': 500 + } + ) + const download = await sampleComfyModelDownload(container, true, startedAt) + expect(download).to.include({ + downloadedBytes: 1500, + totalBytes: 4000, + percent: 38, + filesComplete: 1, + filesInFlight: 1, + filesTotal: 3 + }) + expect(isModelDownloadComplete(download)).to.equal(false) + }) + + it('is complete once every listed file is finished', async () => { + const a = `${HF}/complete-test/repo/resolve/main/a.safetensors` + hubSizes({ [a]: 1000 }) + const container = fakeContainer(`vae\t${a}\n`, new Date(startedAt), { + '/tmp/comfy/models/vae/a.safetensors': 1000 + }) + const download = await sampleComfyModelDownload(container, false, startedAt) + expect(isModelDownloadComplete(download)).to.equal(true) + }) + + it('ignores a list left in the bucket by an earlier launch', async () => { + const a = `${HF}/stale-test/repo/resolve/main/a.safetensors` + const container = fakeContainer(`vae\t${a}\n`, new Date(startedAt - 60_000), {}) + expect(await sampleComfyModelDownload(container, true, startedAt)).to.equal(null) + }) + + it('reports no total when a size is unknown, and nothing when there is no list', async () => { + const a = `${HF}/unknown-test/repo/resolve/main/a.safetensors` + globalThis.fetch = (() => + Promise.resolve(new Response(null, { status: 401 }))) as typeof fetch + const container = fakeContainer(`vae\t${a}\n`, new Date(startedAt), {}) + const download = await sampleComfyModelDownload(container, false, startedAt) + expect(download?.totalBytes).to.equal(undefined) + expect(download?.filesTotal).to.equal(1) + + const empty = { getArchive: () => Promise.reject(new Error('no such file')) } as any + expect(await sampleComfyModelDownload(empty, false, startedAt)).to.equal(null) + }) +}) From 5cebf0c26ff06188818ebd31090ff091e1f1cf29 Mon Sep 17 00:00:00 2001 From: andreip135 <338078672+andreip135@users.noreply.github.com> Date: Thu, 8 Oct 2026 17:29:21 +0300 Subject: [PATCH 2/8] fix automated code review --- src/components/c2d/comfyDownload.ts | 36 ++++++++----- src/test/unit/service/comfyDownload.test.ts | 60 +++++++++++++++++---- 2 files changed, 73 insertions(+), 23 deletions(-) diff --git a/src/components/c2d/comfyDownload.ts b/src/components/c2d/comfyDownload.ts index 04b760a49..60cdc0771 100644 --- a/src/components/c2d/comfyDownload.ts +++ b/src/components/c2d/comfyDownload.ts @@ -57,23 +57,31 @@ function isHubFileUrl(value: string): boolean { ) } -export function parseComfyModelList(text: string, base: string): ComfyModelEntry[] { +/** + * The list's destinations, capped at MAX_ENTRIES. `truncated` says the list names more than that, + * so the entries are not the whole download. + */ +export function parseComfyModelList( + text: string, + base: string +): { entries: ComfyModelEntry[]; truncated: boolean } { const entries = new Map() for (const line of text.split('\n')) { - if (entries.size >= MAX_ENTRIES) { - break - } const [dir, url] = line.trim().split('\t') if (!MODEL_DIR.test(dir) || !url || !isHubFileUrl(url)) { continue } // Same destination the script derives: `$MODELS/$sub/$(basename "$url")`. const path = `${base}/models/${dir}/${url.slice(url.lastIndexOf('/') + 1)}` - if (!entries.has(path)) { - entries.set(path, { url, path }) + if (entries.has(path)) { + continue + } + if (entries.size >= MAX_ENTRIES) { + return { entries: [...entries.values()], truncated: true } } + entries.set(path, { url, path }) } - return [...entries.values()] + return { entries: [...entries.values()], truncated: false } } /** One small file from the container, via Docker's archive endpoint. Null when absent or too big. */ @@ -145,13 +153,17 @@ export async function sampleComfyModelDownload( if (!list || list.mtime < startedAt - 1000) { return null } - const entries = parseComfyModelList(list.text, base) + const { entries, truncated } = parseComfyModelList(list.text, base) if (entries.length === 0) { return null } - const sizes = await Promise.all( - entries.map(({ url }) => cachedHubLookup(url, () => lookupHubFileBytes(url))) - ) + // A truncated list is not the whole download: report bytes only, never a total or a file count + // that would declare it complete early. + const sizes: (number | null)[] = truncated + ? [null] + : await Promise.all( + entries.map(({ url }) => cachedHubLookup(url, () => lookupHubFileBytes(url))) + ) let downloadedBytes = 0 let files = 0 let inFlight = 0 @@ -174,6 +186,6 @@ export async function sampleComfyModelDownload( : null return { ...buildModelDownload({ downloadedBytes, files, inFlight }, totalBytes, null), - filesTotal: entries.length + ...(truncated ? {} : { filesTotal: entries.length }) } } diff --git a/src/test/unit/service/comfyDownload.test.ts b/src/test/unit/service/comfyDownload.test.ts index da7abc5f8..4c64d40e0 100644 --- a/src/test/unit/service/comfyDownload.test.ts +++ b/src/test/unit/service/comfyDownload.test.ts @@ -30,16 +30,31 @@ describe('parseComfyModelList', () => { `loras\t${HF}/org/repo/resolve/main/b.safetensors`, '' ].join('\n') - expect(parseComfyModelList(text, '/tmp/comfy')).to.deep.equal([ - { - url: `${HF}/org/repo/resolve/main/split_files/vae/a.safetensors`, - path: '/tmp/comfy/models/vae/a.safetensors' - }, - { - url: `${HF}/org/repo/resolve/main/b.safetensors`, - path: '/tmp/comfy/models/loras/b.safetensors' - } - ]) + expect(parseComfyModelList(text, '/tmp/comfy')).to.deep.equal({ + entries: [ + { + url: `${HF}/org/repo/resolve/main/split_files/vae/a.safetensors`, + path: '/tmp/comfy/models/vae/a.safetensors' + }, + { + url: `${HF}/org/repo/resolve/main/b.safetensors`, + path: '/tmp/comfy/models/loras/b.safetensors' + } + ], + truncated: false + }) + }) + + it('caps the list and flags one that names more files than the cap', () => { + const line = (i: number) => `vae\t${HF}/org/repo/resolve/main/f${i}.safetensors` + const capped = Array.from({ length: 64 }, (_, i) => line(i)) + const exact = parseComfyModelList([...capped, line(0)].join('\n'), '/tmp/comfy') + expect(exact.entries).to.have.length(64) + expect(exact.truncated).to.equal(false) + + const over = parseComfyModelList([...capped, line(64)].join('\n'), '/tmp/comfy') + expect(over.entries).to.have.length(64) + expect(over.truncated).to.equal(true) }) it('drops anything that is not a plain Hugging Face file URL or a safe directory', () => { @@ -53,7 +68,7 @@ describe('parseComfyModelList', () => { `vae\t${HF}/org/repo/resolve/main/a.bin`, 'vae' ].join('\n') - expect(parseComfyModelList(text, '/tmp/comfy')).to.deep.equal([]) + expect(parseComfyModelList(text, '/tmp/comfy').entries).to.deep.equal([]) }) }) @@ -149,4 +164,27 @@ describe('sampleComfyModelDownload', () => { const empty = { getArchive: () => Promise.reject(new Error('no such file')) } as any expect(await sampleComfyModelDownload(empty, false, startedAt)).to.equal(null) }) + + it('reports bytes only for a list longer than the cap, so it never completes early', async () => { + const url = (i: number) => `${HF}/overflow-test/repo/resolve/main/f${i}.safetensors` + const list = Array.from({ length: 65 }, (_, i) => `vae\t${url(i)}`).join('\n') + const files = Object.fromEntries( + Array.from({ length: 64 }, (_, i) => [ + `/tmp/comfy/models/vae/f${i}.safetensors`, + 10 + ]) + ) + let lookups = 0 + globalThis.fetch = (() => { + lookups++ + return Promise.resolve(new Response(null, { status: 302 })) + }) as typeof fetch + const container = fakeContainer(list, new Date(startedAt), files) + const download = await sampleComfyModelDownload(container, false, startedAt) + expect(download).to.include({ downloadedBytes: 640, filesComplete: 64 }) + expect(download?.filesTotal).to.equal(undefined) + expect(download?.totalBytes).to.equal(undefined) + expect(isModelDownloadComplete(download)).to.equal(false) + expect(lookups).to.equal(0) + }) }) From 56abb3b90960edc0ad93bde4ec6d0f64eb0b938e Mon Sep 17 00:00:00 2001 From: andreip135 <338078672+andreip135@users.noreply.github.com> Date: Thu, 8 Oct 2026 18:29:16 +0300 Subject: [PATCH 3/8] handle template manifests --- docs/services.md | 29 ++++ src/components/c2d/comfyDownload.ts | 84 ++------- src/components/c2d/compute_engine_docker.ts | 105 +++++++++-- src/components/c2d/downloadManifest.ts | 164 ++++++++++++++++++ src/components/c2d/serviceEngines.ts | 21 ++- src/components/database/C2DDatabase.ts | 2 +- src/components/database/sqliteCompute.ts | 5 +- src/test/unit/service/comfyDownload.test.ts | 2 +- .../unit/service/downloadManifest.test.ts | 134 ++++++++++++++ .../unit/service/serviceReadiness.test.ts | 19 +- .../service/unprobedServiceDownload.test.ts | 132 ++++++++++++++ 11 files changed, 610 insertions(+), 87 deletions(-) create mode 100644 src/components/c2d/downloadManifest.ts create mode 100644 src/test/unit/service/downloadManifest.test.ts create mode 100644 src/test/unit/service/unprobedServiceDownload.test.ts diff --git a/docs/services.md b/docs/services.md index 53ae7a4f5..444cacd96 100644 --- a/docs/services.md +++ b/docs/services.md @@ -153,6 +153,35 @@ redeploy) cannot run conflicting operations on the same service. Leases are hear every 30 s while the operation runs; a lease not refreshed for 2 minutes belongs to a crashed process and is stolen automatically, so no manual cleanup is ever needed. +## Readiness and startup progress + +`Running` only means the container started. For workloads it recognizes, the node also reports +whether the service can serve requests and how far its startup downloads have got, on the job +returned by `SERVICE_GET_STATUS`: + +- `readiness` — `waiting`, `ready` or `failing`, from a request the node makes to the workload + itself: vLLM `GET /v1/models` on 8000, llama.cpp `GET /health` on 8080, ComfyUI + `GET /system_stats` on 8188. A profile applies only when the service publishes that port, so an + app that runs an engine behind its own port reports no readiness and `Running` stays the signal. +- `imagePull` — byte progress of the image pull. +- `modelDownload` — bytes downloaded against the total, and `filesComplete` / `filesTotal` when the + download is listed up front. + +`modelDownload` comes from the launch script's download manifest when it writes one: a +tab-separated `/tmp/ocean/downloads.tsv` in the container, one ` []` +line per item, written before the first download. A `file` entry is complete once `` exists +and in flight while `.part` does; a `dir` entry is a folder filled under names of its own +(huggingface_hub's snapshot cache), measured with `du -sb`. A `bytes` of `0` means unknown. Lines may +be appended while the script runs; at most 64 entries are read. Without a manifest, the node falls +back to the engine's own Hugging Face cache (vLLM, llama.cpp) or the `.models.tsv` list older +ComfyUI bundle scripts keep. + +A service with no engine profile (an app that runs its engine behind its own port) still reports +`modelDownload` from its manifest, with no `readiness`. With nothing to say when such a service is +ready, the node checks it only within 30 minutes of the container starting unless a manifest has +appeared, and stops once everything listed is complete. All of it is best-effort and owner-only: a +missing or wrong manifest only changes what that service's own progress report says. + ## Service outputs A service writes its results to `/data/outputs`. With an `outputBucketId` that folder is the diff --git a/src/components/c2d/comfyDownload.ts b/src/components/c2d/comfyDownload.ts index 60cdc0771..4c7d9a6b8 100644 --- a/src/components/c2d/comfyDownload.ts +++ b/src/components/c2d/comfyDownload.ts @@ -1,16 +1,15 @@ import type Dockerode from 'dockerode' -import { pipeline } from 'stream' -import * as tarStream from 'tar-stream' import type { ServiceModelDownload } from '../../@types/C2D/ServiceOnDemand.js' +import { cachedHubLookup, HF_TIMEOUT_MS } from './modelDownload.js' import { - buildModelDownload, - cachedHubLookup, - fileSize, - HF_TIMEOUT_MS -} from './modelDownload.js' + MAX_DOWNLOAD_ENTRIES, + measureDownloads, + readContainerFile +} from './downloadManifest.js' /** - * Model-download progress for the ComfyUI bundle templates. + * Model-download progress for ComfyUI bundle scripts that predate the download manifest (see + * downloadManifest), read from the list they keep for their own download loop. * * Before fetching anything, a bundle's script lists the weights its workflows need in * `/.models.tsv`, one `\t` line each. It then downloads them with curl to @@ -22,8 +21,6 @@ import { */ const MODEL_LIST_FILE = '.models.tsv' -const MAX_LIST_BYTES = 64 * 1024 -const MAX_ENTRIES = 64 const MODEL_DIR = /^[a-z0-9_]{1,32}$/ const URL_SEGMENT = /^[\w.-]+$/ @@ -58,7 +55,7 @@ function isHubFileUrl(value: string): boolean { } /** - * The list's destinations, capped at MAX_ENTRIES. `truncated` says the list names more than that, + * The list's destinations, capped at MAX_DOWNLOAD_ENTRIES. `truncated` says the list names more than that, * so the entries are not the whole download. */ export function parseComfyModelList( @@ -76,7 +73,7 @@ export function parseComfyModelList( if (entries.has(path)) { continue } - if (entries.size >= MAX_ENTRIES) { + if (entries.size >= MAX_DOWNLOAD_ENTRIES) { return { entries: [...entries.values()], truncated: true } } entries.set(path, { url, path }) @@ -84,34 +81,6 @@ export function parseComfyModelList( return { entries: [...entries.values()], truncated: false } } -/** One small file from the container, via Docker's archive endpoint. Null when absent or too big. */ -async function readContainerFile( - container: Dockerode.Container, - path: string -): Promise<{ text: string; mtime: number } | null> { - try { - const extract = tarStream.extract() - pipeline(await container.getArchive({ path }), extract, () => {}) - for await (const entry of extract) { - const { type, size, mtime } = entry.header - if (type !== 'file' || (size ?? 0) > MAX_LIST_BYTES) { - return null - } - const chunks: Buffer[] = [] - for await (const chunk of entry) { - chunks.push(chunk) - } - return { - text: Buffer.concat(chunks).toString('utf8'), - mtime: mtime?.getTime() ?? 0 - } - } - return null - } catch { - return null - } -} - /** * A file's size from the Hub without downloading it. A file in large-file storage answers with a * redirect to the CDN whose `X-Linked-Size` is the file's size; a small one is served directly. @@ -157,35 +126,16 @@ export async function sampleComfyModelDownload( if (entries.length === 0) { return null } - // A truncated list is not the whole download: report bytes only, never a total or a file count - // that would declare it complete early. + // A truncated list is not the whole download: no sizes, so no total or file count that would + // declare it complete early. const sizes: (number | null)[] = truncated - ? [null] + ? entries.map((): null => null) : await Promise.all( entries.map(({ url }) => cachedHubLookup(url, () => lookupHubFileBytes(url))) ) - let downloadedBytes = 0 - let files = 0 - let inFlight = 0 - for (const { path } of entries) { - // Partial first: a file renamed between the two checks is then still found finished. - const partial = await fileSize(container, `${path}.part`) - if (partial !== null) { - downloadedBytes += partial - inFlight++ - continue - } - const finished = await fileSize(container, path) - if (finished !== null) { - downloadedBytes += finished - files++ - } - } - const totalBytes = sizes.every((size) => size !== null) - ? sizes.reduce((sum, size) => sum + size, 0) - : null - return { - ...buildModelDownload({ downloadedBytes, files, inFlight }, totalBytes, null), - ...(truncated ? {} : { filesTotal: entries.length }) - } + return await measureDownloads( + container, + entries.map(({ path }, i) => ({ kind: 'file' as const, path, bytes: sizes[i] })), + truncated + ) } diff --git a/src/components/c2d/compute_engine_docker.ts b/src/components/c2d/compute_engine_docker.ts index d4b32171b..d88a2fefe 100755 --- a/src/components/c2d/compute_engine_docker.ts +++ b/src/components/c2d/compute_engine_docker.ts @@ -100,6 +100,7 @@ import { ModelDownloadSampler } from './modelDownload.js' import { sampleComfyModelDownload } from './comfyDownload.js' +import { MANIFEST_DISCOVERY_MS, sampleDownloadManifest } from './downloadManifest.js' import type { DockerMountObject } from '../../@types/PersistentStorage.js' import { resolveServiceImage } from './serviceResourceMatching.js' import { @@ -216,6 +217,8 @@ export class C2DEngineDocker extends C2DEngine { private serviceProbeUrls: Map = new Map() // Per-service model-download sampling state (which cache files to stat, when to re-list them). private modelDownloadSampler = new ModelDownloadSampler() + // serviceId -> when a service the node cannot probe was last checked for a download manifest. + private manifestCheckedAt: Map = new Map() // serviceId -> its readiness probe (+ model-download sample) still running in the background. // The probe is launched fire-and-forget so a slow engine, Docker daemon or Hub never holds up // InternalLoop; this keeps one probe per service at a time and lets stop() drain them. @@ -3079,13 +3082,16 @@ export class C2DEngineDocker extends C2DEngine { return null } + // With a TTY the output arrives as plain text. Without one, Docker frames it with 8-byte + // headers whose length byte can read as an ASCII digit and get prepended to the size. const exec = await container.exec({ Cmd: ['du', '-sb', path], AttachStdout: true, - AttachStderr: true + AttachStderr: true, + Tty: true }) - const stream = await exec.start({ Detach: false, Tty: false }) + const stream = await exec.start({ Detach: false, Tty: true }) const chunks: Buffer[] = [] for await (const chunk of stream) { @@ -3094,7 +3100,8 @@ export class C2DEngineDocker extends C2DEngine { const output = Buffer.concat(chunks).toString() - const match = output.match(/(\d+)\s/) + // `du -sb` prints "\t"; anything else (an error) is not a size. + const match = output.match(/^\s*(\d+)\s/) return match ? parseInt(match[1], 10) : null } catch (error) { CORE_LOGGER.error( @@ -4645,6 +4652,11 @@ export class C2DEngineDocker extends C2DEngine { this.serviceProbeUrls.delete(serviceId) } } + for (const serviceId of this.manifestCheckedAt.keys()) { + if (!running.has(serviceId)) { + this.manifestCheckedAt.delete(serviceId) + } + } return runningOnly } @@ -4717,7 +4729,10 @@ export class C2DEngineDocker extends C2DEngine { ): Promise { try { const engine = resolveServiceEngine(job) - if (!engine) return + if (!engine) { + await this.sampleUnprobedServiceDownload(job, details) + return + } if (this.serviceOpsInFlight.has(job.serviceId)) return const now = Date.now() @@ -4846,7 +4861,69 @@ export class C2DEngineDocker extends C2DEngine { } /** - * How much of its model the container has pulled down, read from the engine's own cache. + * Download progress for a service the node cannot probe: only what the launch script's download + * manifest says, with no readiness, so Running stays the signal that the service is usable. + * + * Throttled like the probe. A service with no manifest MANIFEST_DISCOVERY_MS after its container + * started is not checked again, and sampling stops once the listed download is complete — with no + * readiness to end it, it would otherwise run for the whole session. + */ + private async sampleUnprobedServiceDownload( + job: ServiceJob, + details: Dockerode.ContainerInspectInfo + ): Promise { + if ( + this.serviceOpsInFlight.has(job.serviceId) || + isModelDownloadComplete(job.modelDownload) + ) { + return + } + const now = Date.now() + const startedAt = Date.parse(details.State?.StartedAt ?? '') + if ( + !job.modelDownload && + Number.isFinite(startedAt) && + now - startedAt > MANIFEST_DISCOVERY_MS + ) { + return + } + if ( + now - (this.manifestCheckedAt.get(job.serviceId) ?? 0) < + PROBE_PERIOD_SECONDS * 1000 + ) { + return + } + this.manifestCheckedAt.set(job.serviceId, now) + const modelDownload = await sampleDownloadManifest( + this.docker.getContainer(job.containerId), + (path) => this.getContainerDiskUsage(job.containerId, path) + ) + if (!modelDownload) { + return + } + // Same cross-process guard as the readiness write: a lifecycle transition must win. + if ( + this.serviceOpsInFlight.has(job.serviceId) || + (await this.db.isServiceLocked(job.serviceId, SERVICE_LOCK_STALE_MS)) + ) { + return + } + await this.db.updateServiceJobReadiness( + job.serviceId, + { + owner: job.owner, + clusterHash: job.clusterHash, + status: ServiceStatusNumber.Running, + containerId: job.containerId + }, + undefined, + modelDownload + ) + } + + /** + * How much of its model the container has pulled down: from the launch script's download + * manifest when it writes one, otherwise from the engine's own cache. * * This is the wait the image pull does not cover: the image is fetched once per node and cached * forever after, while the weights are fetched by the engine on EVERY container start, after it @@ -4861,14 +4938,20 @@ export class C2DEngineDocker extends C2DEngine { engine: ServiceEngineProfile, startedAt: number ): Promise { - if (!engine.modelCachePath && !engine.comfyModelList) return undefined - // Downloaded: what remains is the engine loading weights, which the cache says nothing about. - // Returning nothing keeps the stored 100% record as it is. - if (isModelDownloadComplete(job.modelDownload)) { - return undefined - } try { const container = this.docker.getContainer(job.containerId) + // The launch script's own manifest, when it writes one (see downloadManifest). Read before + // the completion check below: a script can append to it after its first downloads finish. + const fromManifest = await sampleDownloadManifest(container, (path) => + this.getContainerDiskUsage(job.containerId, path) + ) + if (fromManifest) return fromManifest + if (!engine.modelCachePath && !engine.comfyModelList) return undefined + // Downloaded: what remains is the engine loading weights, which the cache says nothing about. + // Returning nothing keeps the stored 100% record as it is. + if (isModelDownloadComplete(job.modelDownload)) { + return undefined + } if (engine.comfyModelList) { const download = await sampleComfyModelDownload( container, diff --git a/src/components/c2d/downloadManifest.ts b/src/components/c2d/downloadManifest.ts new file mode 100644 index 000000000..62ca62113 --- /dev/null +++ b/src/components/c2d/downloadManifest.ts @@ -0,0 +1,164 @@ +import type Dockerode from 'dockerode' +import { pipeline } from 'stream' +import * as tarStream from 'tar-stream' +import type { ServiceModelDownload } from '../../@types/C2D/ServiceOnDemand.js' +import { buildModelDownload, fileSize } from './modelDownload.js' + +/** + * Download progress from a manifest the service's own launch script writes. + * + * Before its first download, a template script lists what it will fetch in MANIFEST_PATH, one + * `\t\t[\t]` line per item: + * + * file one file, written to `.part` and renamed to `` once complete + * dir a folder a tool fills under names of its own (huggingface_hub's snapshot cache), + * measured with `du` and complete once it holds `` + * + * `bytes` is what the item will weigh, or 0 when the script could not tell. A script may append + * lines while it runs, so the file is re-read on every sample. It lives in the container's own + * filesystem, so a new container never sees an earlier one's. Like the script that writes it, the + * manifest is untrusted: it only ever decides what this service's own progress report says. + */ +export const MANIFEST_PATH = '/tmp/ocean/downloads.tsv' +export const MAX_DOWNLOAD_ENTRIES = 64 +// How long after its container starts a service the node cannot probe is checked for a manifest. +// A script writes it before its first download; setup steps before that can take minutes. +export const MANIFEST_DISCOVERY_MS = 30 * 60 * 1000 +const MAX_READ_BYTES = 64 * 1024 +const BYTES = /^\d{1,15}$/ + +export interface DownloadEntry { + kind: 'file' | 'dir' + path: string + bytes: number | null +} + +/** Bytes held under a folder in the container, or null when it is not there yet. */ +export type DirectoryBytes = (path: string) => Promise + +/** One small file from the container, via Docker's archive endpoint. Null when absent or too big. */ +export async function readContainerFile( + container: Dockerode.Container, + path: string +): Promise<{ text: string; mtime: number } | null> { + try { + const extract = tarStream.extract() + pipeline(await container.getArchive({ path }), extract, () => {}) + for await (const entry of extract) { + const { type, size, mtime } = entry.header + if (type !== 'file' || (size ?? 0) > MAX_READ_BYTES) { + return null + } + const chunks: Buffer[] = [] + for await (const chunk of entry) { + chunks.push(chunk) + } + return { + text: Buffer.concat(chunks).toString('utf8'), + mtime: mtime?.getTime() ?? 0 + } + } + return null + } catch { + return null + } +} + +/** + * The manifest's entries, capped at MAX_DOWNLOAD_ENTRIES. `truncated` says it names more than + * that, so the entries are not the whole download. + */ +export function parseDownloadManifest(text: string): { + entries: DownloadEntry[] + truncated: boolean +} { + const entries = new Map() + for (const line of text.split('\n')) { + const [kind, bytes, path] = line.trim().split('\t') + if ( + (kind !== 'file' && kind !== 'dir') || + !BYTES.test(bytes ?? '') || + !path?.startsWith('/') || + path.includes('\0') || + entries.has(path) + ) { + continue + } + if (entries.size >= MAX_DOWNLOAD_ENTRIES) { + return { entries: [...entries.values()], truncated: true } + } + entries.set(path, { kind, path, bytes: Number(bytes) || null }) + } + return { entries: [...entries.values()], truncated: false } +} + +/** + * Bytes on disk against the sizes the entries declare. The total and the file count are reported + * only when they are the whole download, so a truncated list never reads as complete. + */ +export async function measureDownloads( + container: Dockerode.Container, + entries: DownloadEntry[], + truncated: boolean, + directoryBytes?: DirectoryBytes +): Promise { + let downloadedBytes = 0 + let files = 0 + let inFlight = 0 + for (const { kind, path, bytes } of entries) { + if (kind === 'dir') { + const held = directoryBytes ? await directoryBytes(path) : null + if (held === null) { + continue + } + // Capped: a reused cache can hold more than this launch asked for. + downloadedBytes += bytes ? Math.min(held, bytes) : held + if (bytes && held >= bytes) { + files++ + } else { + inFlight++ + } + continue + } + // Partial first: a file renamed between the two checks is then still found finished. + const partial = await fileSize(container, `${path}.part`) + if (partial !== null) { + downloadedBytes += partial + inFlight++ + continue + } + const finished = await fileSize(container, path) + if (finished !== null) { + downloadedBytes += finished + files++ + } + } + const totalBytes = + !truncated && entries.every(({ bytes }) => bytes !== null) + ? entries.reduce((sum, { bytes }) => sum + bytes, 0) + : null + return { + ...buildModelDownload({ downloadedBytes, files, inFlight }, totalBytes, null), + ...(truncated ? {} : { filesTotal: entries.length }) + } +} + +/** Progress from the launch script's manifest. Null when the script writes none. Never throws. */ +export async function sampleDownloadManifest( + container: Dockerode.Container, + directoryBytes: DirectoryBytes +): Promise { + const manifest = await readContainerFile(container, MANIFEST_PATH) + if (!manifest) { + return null + } + const { entries, truncated } = parseDownloadManifest(manifest.text) + if (entries.length === 0) { + return null + } + try { + return await measureDownloads(container, entries, truncated, directoryBytes) + } catch { + return null + } +} diff --git a/src/components/c2d/serviceEngines.ts b/src/components/c2d/serviceEngines.ts index 3f5a70c75..f752cf68d 100644 --- a/src/components/c2d/serviceEngines.ts +++ b/src/components/c2d/serviceEngines.ts @@ -175,11 +175,28 @@ function stripImageRef(image: string): string { return colon === -1 ? withoutDigest : withoutDigest.slice(0, colon) } -/** The engine profile for a service, or null when the node does not recognize the image. */ +/** + * The engine profile for a service, or null when the node does not recognize the image — or when + * the service does not publish the engine's port. An app that runs the engine behind itself (vLLM + * on loopback, the app on its own port) is not the engine's to call ready: probing the app's port + * for the engine's path would never succeed. + */ export function resolveServiceEngine(job: ServiceJob): ServiceEngineProfile | null { const image = stripImageRef((job.image || job.containerImage || '').trim()) if (!image) { return null } - return SERVICE_ENGINE_PROFILES.find((profile) => profile.matchesImage(image)) ?? null + const profile = SERVICE_ENGINE_PROFILES.find((candidate) => + candidate.matchesImage(image) + ) + if (!profile) { + return null + } + if ( + profile.probe.port !== undefined && + !job.exposedPorts?.includes(profile.probe.port) + ) { + return null + } + return profile } diff --git a/src/components/database/C2DDatabase.ts b/src/components/database/C2DDatabase.ts index 9fbf0379d..72a43c049 100755 --- a/src/components/database/C2DDatabase.ts +++ b/src/components/database/C2DDatabase.ts @@ -120,7 +120,7 @@ export class C2DDatabase extends AbstractDatabase { status: number containerId: string }, - readiness: ServiceReadiness, + readiness: ServiceReadiness | undefined, modelDownload?: ServiceModelDownload ): Promise { return await this.provider.updateServiceJobReadiness( diff --git a/src/components/database/sqliteCompute.ts b/src/components/database/sqliteCompute.ts index 90890cb36..b034fb111 100644 --- a/src/components/database/sqliteCompute.ts +++ b/src/components/database/sqliteCompute.ts @@ -409,6 +409,7 @@ export class SQLiteCompute implements ComputeDatabaseProvider { // Same guarantees as updateServiceJobMetrics, for the readiness probe result: it is sampled from // the same lease-free background loop, so it must never overwrite a lifecycle transition either. + // `readiness` is absent for a service the node cannot probe, which reports download progress only. // eslint-disable-next-line require-await async updateServiceJobReadiness( serviceId: string, @@ -418,11 +419,11 @@ export class SQLiteCompute implements ComputeDatabaseProvider { status: number containerId: string }, - readiness: ServiceReadiness, + readiness: ServiceReadiness | undefined, modelDownload?: ServiceModelDownload ): Promise { return this.patchServiceJobBody(serviceId, expected, 'readiness', (body) => { - body.readiness = readiness + if (readiness) body.readiness = readiness // Only overwritten when a fresh sample was taken: once the engine is ready the walk stops, // and the last figures stay as the record of what was downloaded. if (modelDownload) body.modelDownload = modelDownload diff --git a/src/test/unit/service/comfyDownload.test.ts b/src/test/unit/service/comfyDownload.test.ts index 4c64d40e0..8c7983884 100644 --- a/src/test/unit/service/comfyDownload.test.ts +++ b/src/test/unit/service/comfyDownload.test.ts @@ -11,7 +11,7 @@ import { const HF = 'https://huggingface.co' describe('ComfyUI engine profile', () => { - const job = (image: string) => ({ image }) as ServiceJob + const job = (image: string) => ({ image, exposedPorts: [8188] }) as ServiceJob it('recognizes the comfyui-boot image under any tag', () => { expect(resolveServiceEngine(job('yanwk/comfyui-boot'))?.id).to.equal('comfyui') diff --git a/src/test/unit/service/downloadManifest.test.ts b/src/test/unit/service/downloadManifest.test.ts new file mode 100644 index 000000000..2acbb1b7d --- /dev/null +++ b/src/test/unit/service/downloadManifest.test.ts @@ -0,0 +1,134 @@ +import { expect } from 'chai' +import * as tarStream from 'tar-stream' +import { isModelDownloadComplete } from '../../../components/c2d/modelDownload.js' +import { + MANIFEST_PATH, + parseDownloadManifest, + sampleDownloadManifest +} from '../../../components/c2d/downloadManifest.js' + +const MODELS = '/data/outputs/comfy/models' +const GEMMA = '/data/outputs/comfy/.cache/huggingface/hub/models--google--gemma-4-31B-it' + +describe('parseDownloadManifest', () => { + it('reads file and dir entries, treating an unknown size as null', () => { + const text = [ + `file\t1000\t${MODELS}/vae/a.safetensors\thttps://huggingface.co/org/repo/resolve/main/a.safetensors`, + `file\t0\t${MODELS}/loras/b.safetensors`, + `dir\t5000\t${GEMMA}`, + '' + ].join('\n') + expect(parseDownloadManifest(text)).to.deep.equal({ + entries: [ + { kind: 'file', path: `${MODELS}/vae/a.safetensors`, bytes: 1000 }, + { kind: 'file', path: `${MODELS}/loras/b.safetensors`, bytes: null }, + { kind: 'dir', path: GEMMA, bytes: 5000 } + ], + truncated: false + }) + }) + + it('drops malformed lines and repeated paths', () => { + const text = [ + `blob\t1\t${MODELS}/a`, + `file\t-1\t${MODELS}/a`, + `file\tlots\t${MODELS}/a`, + 'file\t1\trelative/path', + `file\t1\t${MODELS}/a\u0000b`, + `file\t1\t${MODELS}/a`, + `file\t2\t${MODELS}/a`, + 'file\t1' + ].join('\n') + expect(parseDownloadManifest(text).entries).to.deep.equal([ + { kind: 'file', path: `${MODELS}/a`, bytes: 1 } + ]) + }) + + it('caps the entries and flags a manifest that names more', () => { + const line = (i: number) => `file\t1\t${MODELS}/f${i}` + const capped = Array.from({ length: 64 }, (_, i) => line(i)) + expect(parseDownloadManifest(capped.join('\n')).truncated).to.equal(false) + const over = parseDownloadManifest([...capped, line(64)].join('\n')) + expect(over.entries).to.have.length(64) + expect(over.truncated).to.equal(true) + }) +}) + +describe('sampleDownloadManifest', () => { + // A container double: `getArchive()` serves the manifest as a tar, `infoArchive()` answers stats. + function fakeContainer(manifest: string | null, files: Record) { + return { + getArchive: ({ path }: { path: string }) => { + if (manifest === null || path !== MANIFEST_PATH) { + return Promise.reject(new Error('no such file')) + } + const pack = tarStream.pack() + pack.entry({ name: 'downloads.tsv' }, manifest) + pack.finalize() + return Promise.resolve(pack) + }, + infoArchive: ({ path }: { path: string }) => { + if (files[path] === undefined) { + return Promise.reject(new Error('not found')) + } + const header = Buffer.from(JSON.stringify({ size: files[path] })).toString( + 'base64' + ) + return Promise.resolve({ headers: { 'x-docker-container-path-stat': header } }) + } + } as any + } + + const manifest = [ + `file\t1000\t${MODELS}/vae/a.safetensors`, + `file\t2000\t${MODELS}/loras/b.safetensors`, + `dir\t7000\t${GEMMA}` + ].join('\n') + + it('measures files and folders against the sizes the script declared', async () => { + const container = fakeContainer(manifest, { + [`${MODELS}/vae/a.safetensors`]: 1000, + [`${MODELS}/loras/b.safetensors.part`]: 500 + }) + const download = await sampleDownloadManifest(container, (path) => + Promise.resolve(path === GEMMA ? 3000 : null) + ) + expect(download).to.include({ + downloadedBytes: 4500, + totalBytes: 10000, + percent: 45, + filesComplete: 1, + filesInFlight: 2, + filesTotal: 3 + }) + expect(isModelDownloadComplete(download)).to.equal(false) + }) + + it('is complete once every file is in place and every folder holds its size', async () => { + const container = fakeContainer(manifest, { + [`${MODELS}/vae/a.safetensors`]: 1000, + [`${MODELS}/loras/b.safetensors`]: 2000 + }) + // A reused cache can hold more than declared; it counts only up to the declared size. + const download = await sampleDownloadManifest(container, () => Promise.resolve(9000)) + expect(download).to.include({ + downloadedBytes: 10000, + percent: 100, + filesComplete: 3 + }) + expect(isModelDownloadComplete(download)).to.equal(true) + }) + + it('reports no total when a size is unknown', async () => { + const container = fakeContainer(`file\t0\t${MODELS}/vae/a.safetensors`, {}) + const download = await sampleDownloadManifest(container, () => Promise.resolve(null)) + expect(download?.totalBytes).to.equal(undefined) + expect(download?.filesTotal).to.equal(1) + }) + + it('returns null when the script writes no manifest, or an empty one', async () => { + const none = () => Promise.resolve(null) + expect(await sampleDownloadManifest(fakeContainer(null, {}), none)).to.equal(null) + expect(await sampleDownloadManifest(fakeContainer('', {}), none)).to.equal(null) + }) +}) diff --git a/src/test/unit/service/serviceReadiness.test.ts b/src/test/unit/service/serviceReadiness.test.ts index cf55eabe6..7398f7a4a 100644 --- a/src/test/unit/service/serviceReadiness.test.ts +++ b/src/test/unit/service/serviceReadiness.test.ts @@ -118,7 +118,7 @@ describe('ImagePullTracker', () => { }) describe('resolveServiceEngine', () => { - const job = (image: string) => ({ image }) as ServiceJob + const job = (image: string) => ({ image, exposedPorts: [8000, 8080] }) as ServiceJob it('recognizes the vLLM image, including a registry mirror of it', () => { expect(resolveServiceEngine(job('vllm/vllm-openai'))?.id).to.equal('vllm') @@ -140,7 +140,10 @@ describe('resolveServiceEngine', () => { }) it('falls back to containerImage when image is missing', () => { - const noImage = { containerImage: 'vllm/vllm-openai:v0.28.0' } as ServiceJob + const noImage = { + containerImage: 'vllm/vllm-openai:v0.28.0', + exposedPorts: [8000] + } as ServiceJob expect(resolveServiceEngine(noImage)?.id).to.equal('vllm') }) @@ -150,6 +153,15 @@ describe('resolveServiceEngine', () => { expect(resolveServiceEngine(job(''))).to.equal(null) }) + it('leaves an app that runs the engine behind its own port unrecognized', () => { + // A bundle that publishes only its app (vLLM on loopback inside) has nothing to probe. + const app = { image: 'vllm/vllm-openai', exposedPorts: [9119] } as ServiceJob + expect(resolveServiceEngine(app)).to.equal(null) + expect(resolveServiceEngine({ image: 'vllm/vllm-openai' } as ServiceJob)).to.equal( + null + ) + }) + it('reads the Hugging Face repo id out of the vLLM command', () => { const vllm = resolveServiceEngine(job('vllm/vllm-openai'))! expect( @@ -169,7 +181,8 @@ describe('resolveServiceEngine', () => { }) describe('resolveServiceEngine — llama.cpp', () => { - const job = (image: string, cmd?: string[]) => ({ image, dockerCmd: cmd }) as ServiceJob + const job = (image: string, cmd?: string[]) => + ({ image, dockerCmd: cmd, exposedPorts: [8000, 8080] }) as ServiceJob it('recognizes the llama.cpp image and probes /health', () => { const engine = resolveServiceEngine(job('ghcr.io/ggml-org/llama.cpp')) diff --git a/src/test/unit/service/unprobedServiceDownload.test.ts b/src/test/unit/service/unprobedServiceDownload.test.ts new file mode 100644 index 000000000..b5e7e3f55 --- /dev/null +++ b/src/test/unit/service/unprobedServiceDownload.test.ts @@ -0,0 +1,132 @@ +import { expect } from 'chai' +import sinon from 'sinon' +import { Readable } from 'stream' +import * as tarStream from 'tar-stream' +import { C2DEngineDocker } from '../../../components/c2d/compute_engine_docker.js' +import { ServiceStatusNumber } from '../../../@types/C2D/ServiceOnDemand.js' + +// A vLLM-based app bundle: the engine image, but only the app's own port published, so the node +// has no engine to probe and reports the launch script's download manifest alone. +describe('download progress for a service the node cannot probe', () => { + const CACHE = '/data/outputs/hf/hub/models--Qwen--Qwen3.8-27B-FP8' + + function engineWith(manifest: string | null, folderBytes: number | null) { + const engine: any = Object.create(C2DEngineDocker.prototype) + engine.serviceOpsInFlight = new Set() + engine.manifestCheckedAt = new Map() + engine.getArchive = sinon.stub().callsFake(() => { + if (manifest === null) { + return Promise.reject(new Error('no such file')) + } + const pack = tarStream.pack() + pack.entry({ name: 'downloads.tsv' }, manifest) + pack.finalize() + return Promise.resolve(pack) + }) + engine.docker = { getContainer: () => ({ getArchive: engine.getArchive }) } + engine.getContainerDiskUsage = sinon.stub().resolves(folderBytes) + engine.db = { + isServiceLocked: sinon.stub().resolves(false), + updateServiceJobReadiness: sinon.stub().resolves(true) + } + return engine + } + + const job = (extra: object = {}) => ({ + serviceId: 'svc-1', + owner: '0x0000000000000000000000000000000000000001', + clusterHash: '0xcluster', + containerId: 'container-1', + image: 'vllm/vllm-openai', + exposedPorts: [9119], + status: ServiceStatusNumber.Running, + ...extra + }) + const startedAgo = (ms: number) => ({ + State: { StartedAt: new Date(Date.now() - ms).toISOString() } + }) + + it('records the manifest progress with no readiness', async () => { + const engine = engineWith(`dir\t1000\t${CACHE}\n`, 400) + await engine.probeServiceReadiness(job(), startedAgo(60_000)) + sinon.assert.calledOnce(engine.db.updateServiceJobReadiness) + const [serviceId, expected, readiness, download] = + engine.db.updateServiceJobReadiness.firstCall.args + expect(serviceId).to.equal('svc-1') + expect(expected.status).to.equal(ServiceStatusNumber.Running) + expect(readiness).to.equal(undefined) + expect(download).to.include({ + downloadedBytes: 400, + totalBytes: 1000, + percent: 40, + filesTotal: 1 + }) + sinon.assert.calledWith(engine.getContainerDiskUsage, 'container-1', CACHE) + }) + + it('checks at most once per probe period', async () => { + const engine = engineWith(`dir\t1000\t${CACHE}\n`, 400) + await engine.probeServiceReadiness(job(), startedAgo(60_000)) + await engine.probeServiceReadiness(job(), startedAgo(60_000)) + sinon.assert.calledOnce(engine.getArchive) + }) + + it('stops once the listed download is complete', async () => { + const engine = engineWith(`dir\t1000\t${CACHE}\n`, 1000) + const complete = { + downloadedBytes: 1000, + totalBytes: 1000, + percent: 100, + filesComplete: 1, + filesInFlight: 0, + filesTotal: 1, + updatedAt: Date.now() + } + await engine.probeServiceReadiness( + job({ modelDownload: complete }), + startedAgo(60_000) + ) + sinon.assert.notCalled(engine.getArchive) + sinon.assert.notCalled(engine.db.updateServiceJobReadiness) + }) + + it('stops looking once the discovery window has passed with no manifest', async () => { + const engine = engineWith(null, null) + await engine.probeServiceReadiness(job(), startedAgo(31 * 60 * 1000)) + sinon.assert.notCalled(engine.getArchive) + }) + + it('writes nothing for a service that has no manifest', async () => { + const engine = engineWith(null, null) + await engine.probeServiceReadiness(job(), startedAgo(60_000)) + sinon.assert.calledOnce(engine.getArchive) + sinon.assert.notCalled(engine.db.updateServiceJobReadiness) + }) +}) + +describe('getContainerDiskUsage', () => { + function engineWithOutput(output: string) { + const engine: any = Object.create(C2DEngineDocker.prototype) + const exec = { start: sinon.stub().resolves(Readable.from([Buffer.from(output)])) } + engine.docker = { + getContainer: () => ({ + inspect: sinon.stub().resolves({ State: { Running: true } }), + exec: sinon.stub().resolves(exec) + }) + } + return engine + } + + it('reads the size du prints for a long path', async () => { + const path = '/data/outputs/hf/hub/models--Qwen--Qwen3.8-test' + const engine = engineWithOutput(`5004096\t${path}\r\n`) + expect(await engine.getContainerDiskUsage('container-1', path)).to.equal(5004096) + }) + + it('never reads an error message as a size', async () => { + const engine = engineWithOutput( + "du: cannot access '/data/models--org--model-2 1': No such file or directory\r\n" + ) + expect(await engine.getContainerDiskUsage('container-1', '/data/x')).to.equal(null) + }) +}) From c7269b8de7b9e61a41cdbcd5fb42056faa9bc8af Mon Sep 17 00:00:00 2001 From: andreip135 <338078672+andreip135@users.noreply.github.com> Date: Thu, 8 Oct 2026 18:48:42 +0300 Subject: [PATCH 4/8] add du timeout --- src/components/c2d/compute_engine_docker.ts | 27 +++++++++++++++++-- .../service/unprobedServiceDownload.test.ts | 27 ++++++++++++++++++- 2 files changed, 51 insertions(+), 3 deletions(-) diff --git a/src/components/c2d/compute_engine_docker.ts b/src/components/c2d/compute_engine_docker.ts index d88a2fefe..9df01ad46 100755 --- a/src/components/c2d/compute_engine_docker.ts +++ b/src/components/c2d/compute_engine_docker.ts @@ -147,6 +147,9 @@ const trivyImage = 'aquasec/trivy:0.69.3' // Use pinned versions for safety // doesn't get a live operation's resources torn down under it. const SERVICE_LOCK_HEARTBEAT_MS = 30_000 const SERVICE_LOCK_STALE_MS = 120_000 +// A `du` inside a container that never finishes (a huge or slow folder, or a container that replaced +// the binary) must not hold the loop that measures it, nor the node's shutdown that drains it. +const DISK_USAGE_TIMEOUT_MS = 15_000 // Identifies one engine instance as a service-lock holder across processes. pid alone is // not enough: pids are reused, and one process can host several engine instances. @@ -3094,8 +3097,28 @@ export class C2DEngineDocker extends C2DEngine { const stream = await exec.start({ Detach: false, Tty: true }) const chunks: Buffer[] = [] - for await (const chunk of stream) { - chunks.push(chunk as Buffer) + const read = (async () => { + for await (const chunk of stream) { + chunks.push(chunk as Buffer) + } + return true + })() + // Destroying the stream on timeout rejects the read after the race has moved on. + read.catch(() => {}) + let timer: NodeJS.Timeout + const timedOut = new Promise((resolve) => { + timer = setTimeout(() => resolve(false), DISK_USAGE_TIMEOUT_MS) + }) + try { + if (!(await Promise.race([read, timedOut]))) { + stream.destroy() + CORE_LOGGER.warn( + `du ${path} in ${containerName} did not finish within ${DISK_USAGE_TIMEOUT_MS / 1000}s` + ) + return null + } + } finally { + clearTimeout(timer) } const output = Buffer.concat(chunks).toString() diff --git a/src/test/unit/service/unprobedServiceDownload.test.ts b/src/test/unit/service/unprobedServiceDownload.test.ts index b5e7e3f55..a6e7b57dd 100644 --- a/src/test/unit/service/unprobedServiceDownload.test.ts +++ b/src/test/unit/service/unprobedServiceDownload.test.ts @@ -1,6 +1,6 @@ import { expect } from 'chai' import sinon from 'sinon' -import { Readable } from 'stream' +import { PassThrough, Readable } from 'stream' import * as tarStream from 'tar-stream' import { C2DEngineDocker } from '../../../components/c2d/compute_engine_docker.js' import { ServiceStatusNumber } from '../../../@types/C2D/ServiceOnDemand.js' @@ -130,3 +130,28 @@ describe('getContainerDiskUsage', () => { expect(await engine.getContainerDiskUsage('container-1', '/data/x')).to.equal(null) }) }) + +describe('getContainerDiskUsage timeout', () => { + let clock: sinon.SinonFakeTimers + beforeEach(() => { + clock = sinon.useFakeTimers({ toFake: ['setTimeout', 'clearTimeout'] }) + }) + afterEach(() => { + clock.restore() + }) + + it('gives up on a du that never finishes, instead of hanging', async () => { + const stream = new PassThrough() + const engine: any = Object.create(C2DEngineDocker.prototype) + engine.docker = { + getContainer: () => ({ + inspect: sinon.stub().resolves({ State: { Running: true } }), + exec: sinon.stub().resolves({ start: sinon.stub().resolves(stream) }) + }) + } + const result = engine.getContainerDiskUsage('container-1', '/data/hung') + await clock.tickAsync(15_000) + expect(await result).to.equal(null) + expect(stream.destroyed).to.equal(true) + }) +}) From 3d7d769c58b71dab168a5330c21032cc20cfd208 Mon Sep 17 00:00:00 2001 From: andreip135 <338078672+andreip135@users.noreply.github.com> Date: Thu, 8 Oct 2026 18:58:58 +0300 Subject: [PATCH 5/8] warn when a service is still not ready 15min after start, replace .models.tsv approach with download manifest --- docs/services.md | 5 +- src/components/c2d/comfyDownload.ts | 141 ------------- src/components/c2d/compute_engine_docker.ts | 56 ++++-- src/components/c2d/downloadManifest.ts | 15 +- src/components/c2d/modelDownload.ts | 4 +- src/components/c2d/serviceEngines.ts | 5 +- src/components/c2d/serviceReadiness.ts | 3 + src/test/unit/service/comfyDownload.test.ts | 190 ------------------ .../unit/service/readinessWarning.test.ts | 73 +++++++ .../unit/service/serviceReadiness.test.ts | 12 ++ 10 files changed, 134 insertions(+), 370 deletions(-) delete mode 100644 src/components/c2d/comfyDownload.ts delete mode 100644 src/test/unit/service/comfyDownload.test.ts create mode 100644 src/test/unit/service/readinessWarning.test.ts diff --git a/docs/services.md b/docs/services.md index 444cacd96..848e9d431 100644 --- a/docs/services.md +++ b/docs/services.md @@ -163,6 +163,8 @@ returned by `SERVICE_GET_STATUS`: itself: vLLM `GET /v1/models` on 8000, llama.cpp `GET /health` on 8080, ComfyUI `GET /system_stats` on 8188. A profile applies only when the service publishes that port, so an app that runs an engine behind its own port reports no readiness and `Running` stays the signal. + A service still `waiting` 15 minutes after its container started is logged once as a node warning + with the last probe error — usually a node that cannot reach its containers. - `imagePull` — byte progress of the image pull. - `modelDownload` — bytes downloaded against the total, and `filesComplete` / `filesTotal` when the download is listed up front. @@ -173,8 +175,7 @@ line per item, written before the first download. A `file` entry is complete onc and in flight while `.part` does; a `dir` entry is a folder filled under names of its own (huggingface_hub's snapshot cache), measured with `du -sb`. A `bytes` of `0` means unknown. Lines may be appended while the script runs; at most 64 entries are read. Without a manifest, the node falls -back to the engine's own Hugging Face cache (vLLM, llama.cpp) or the `.models.tsv` list older -ComfyUI bundle scripts keep. +back to the engine's own Hugging Face cache (vLLM, llama.cpp). A service with no engine profile (an app that runs its engine behind its own port) still reports `modelDownload` from its manifest, with no `readiness`. With nothing to say when such a service is diff --git a/src/components/c2d/comfyDownload.ts b/src/components/c2d/comfyDownload.ts deleted file mode 100644 index 4c7d9a6b8..000000000 --- a/src/components/c2d/comfyDownload.ts +++ /dev/null @@ -1,141 +0,0 @@ -import type Dockerode from 'dockerode' -import type { ServiceModelDownload } from '../../@types/C2D/ServiceOnDemand.js' -import { cachedHubLookup, HF_TIMEOUT_MS } from './modelDownload.js' -import { - MAX_DOWNLOAD_ENTRIES, - measureDownloads, - readContainerFile -} from './downloadManifest.js' - -/** - * Model-download progress for ComfyUI bundle scripts that predate the download manifest (see - * downloadManifest), read from the list they keep for their own download loop. - * - * Before fetching anything, a bundle's script lists the weights its workflows need in - * `/.models.tsv`, one `\t` line each. It then downloads them with curl to - * `/models//.part`, renaming each to `` once its size checks out. `` is - * `/data/outputs/comfy` when a bucket is mounted, `/tmp/comfy` otherwise. - * - * The script is client-supplied, so the list is untrusted: only plain Hugging Face file URLs are - * kept, and those are the only addresses the node ever contacts on its behalf. - */ - -const MODEL_LIST_FILE = '.models.tsv' -const MODEL_DIR = /^[a-z0-9_]{1,32}$/ -const URL_SEGMENT = /^[\w.-]+$/ - -export interface ComfyModelEntry { - url: string - path: string -} - -export function comfyBaseDir(hasBucket: boolean): string { - return hasBucket ? '/data/outputs/comfy' : '/tmp/comfy' -} - -/** `https://huggingface.co///resolve//.safetensors`, exactly. */ -function isHubFileUrl(value: string): boolean { - let url: URL - try { - url = new URL(value) - } catch { - return false - } - // Rebuilt from the parsed path, so a query, credentials, a port or a normalized `..` all fail. - if (value !== `https://huggingface.co${url.pathname}`) { - return false - } - const segments = url.pathname.split('/').slice(1) - return ( - segments.length >= 5 && - segments[2] === 'resolve' && - segments.every((segment) => URL_SEGMENT.test(segment)) && - url.pathname.endsWith('.safetensors') - ) -} - -/** - * The list's destinations, capped at MAX_DOWNLOAD_ENTRIES. `truncated` says the list names more than that, - * so the entries are not the whole download. - */ -export function parseComfyModelList( - text: string, - base: string -): { entries: ComfyModelEntry[]; truncated: boolean } { - const entries = new Map() - for (const line of text.split('\n')) { - const [dir, url] = line.trim().split('\t') - if (!MODEL_DIR.test(dir) || !url || !isHubFileUrl(url)) { - continue - } - // Same destination the script derives: `$MODELS/$sub/$(basename "$url")`. - const path = `${base}/models/${dir}/${url.slice(url.lastIndexOf('/') + 1)}` - if (entries.has(path)) { - continue - } - if (entries.size >= MAX_DOWNLOAD_ENTRIES) { - return { entries: [...entries.values()], truncated: true } - } - entries.set(path, { url, path }) - } - return { entries: [...entries.values()], truncated: false } -} - -/** - * A file's size from the Hub without downloading it. A file in large-file storage answers with a - * redirect to the CDN whose `X-Linked-Size` is the file's size; a small one is served directly. - */ -async function lookupHubFileBytes(url: string): Promise { - try { - const response = await fetch(url, { - method: 'HEAD', - redirect: 'manual', - signal: AbortSignal.timeout(HF_TIMEOUT_MS) - }) - if (response.status >= 400) { - // Missing, or gated without a token: asking again will not change the answer. - return [401, 403, 404].includes(response.status) ? null : undefined - } - const size = Number( - response.headers.get('x-linked-size') ?? - (response.ok ? response.headers.get('content-length') : null) - ) - return size > 0 ? size : null - } catch { - return undefined - } -} - -/** - * Bytes on disk against the size of everything the list names. Null until the list exists — and, - * in a reused bucket, until this launch has rewritten the previous one (hence the mtime check). - * Never throws. - */ -export async function sampleComfyModelDownload( - container: Dockerode.Container, - hasBucket: boolean, - startedAt: number -): Promise { - const base = comfyBaseDir(hasBucket) - const list = await readContainerFile(container, `${base}/${MODEL_LIST_FILE}`) - // tar keeps whole seconds, hence the one-second slack. - if (!list || list.mtime < startedAt - 1000) { - return null - } - const { entries, truncated } = parseComfyModelList(list.text, base) - if (entries.length === 0) { - return null - } - // A truncated list is not the whole download: no sizes, so no total or file count that would - // declare it complete early. - const sizes: (number | null)[] = truncated - ? entries.map((): null => null) - : await Promise.all( - entries.map(({ url }) => cachedHubLookup(url, () => lookupHubFileBytes(url))) - ) - return await measureDownloads( - container, - entries.map(({ path }, i) => ({ kind: 'file' as const, path, bytes: sizes[i] })), - truncated - ) -} diff --git a/src/components/c2d/compute_engine_docker.ts b/src/components/c2d/compute_engine_docker.ts index 9df01ad46..50601f4bf 100755 --- a/src/components/c2d/compute_engine_docker.ts +++ b/src/components/c2d/compute_engine_docker.ts @@ -90,7 +90,8 @@ import { runReadinessProbe, PROBE_INITIAL_DELAY_SECONDS, PROBE_PERIOD_SECONDS, - READY_PROBE_PERIOD_SECONDS + READY_PROBE_PERIOD_SECONDS, + READINESS_WARN_AFTER_SECONDS } from './serviceReadiness.js' import { resolveServiceEngine, type ServiceEngineProfile } from './serviceEngines.js' import { @@ -99,7 +100,6 @@ import { isModelDownloadComplete, ModelDownloadSampler } from './modelDownload.js' -import { sampleComfyModelDownload } from './comfyDownload.js' import { MANIFEST_DISCOVERY_MS, sampleDownloadManifest } from './downloadManifest.js' import type { DockerMountObject } from '../../@types/PersistentStorage.js' import { resolveServiceImage } from './serviceResourceMatching.js' @@ -222,6 +222,8 @@ export class C2DEngineDocker extends C2DEngine { private modelDownloadSampler = new ModelDownloadSampler() // serviceId -> when a service the node cannot probe was last checked for a download manifest. private manifestCheckedAt: Map = new Map() + // serviceId -> the container a "still not ready" warning was logged for (once per container). + private readinessWarned: Map = new Map() // serviceId -> its readiness probe (+ model-download sample) still running in the background. // The probe is launched fire-and-forget so a slow engine, Docker daemon or Hub never holds up // InternalLoop; this keeps one probe per service at a time and lets stop() drain them. @@ -4670,14 +4672,15 @@ export class C2DEngineDocker extends C2DEngine { // expired), whichever path ended them. const running = new Set(services.map((svc) => svc.serviceId)) this.modelDownloadSampler.retain(running) - for (const serviceId of this.serviceProbeUrls.keys()) { - if (!running.has(serviceId)) { - this.serviceProbeUrls.delete(serviceId) - } - } - for (const serviceId of this.manifestCheckedAt.keys()) { - if (!running.has(serviceId)) { - this.manifestCheckedAt.delete(serviceId) + for (const state of [ + this.serviceProbeUrls, + this.manifestCheckedAt, + this.readinessWarned + ]) { + for (const serviceId of state.keys()) { + if (!running.has(serviceId)) { + state.delete(serviceId) + } } } return runningOnly @@ -4856,7 +4859,25 @@ export class C2DEngineDocker extends C2DEngine { // when it says nothing about why. const modelDownload = everReady ? undefined - : await this.sampleModelDownload(job, engine, startedAt) + : await this.sampleModelDownload(job, engine) + + // A service that never answers is usually a node that cannot reach its container: say so once, + // where an operator looks, rather than only through users stuck on "warming up". + if ( + readiness.state === 'waiting' && + Number.isFinite(startedAt) && + now - startedAt >= READINESS_WARN_AFTER_SECONDS * 1000 && + this.readinessWarned.get(job.serviceId) !== job.containerId + ) { + this.readinessWarned.set(job.serviceId, job.containerId) + const download = modelDownload ?? job.modelDownload + CORE_LOGGER.warn( + `[readiness] service ${job.serviceId} (${engine.id}) not ready ` + + `${Math.round((now - startedAt) / 60_000)} min after its container started: ` + + `${readiness.lastError ?? `HTTP ${readiness.httpStatus}`} via ${readiness.probedUrl}` + + (download ? `; model download at ${download.percent ?? '?'}%` : '') + ) + } // Same cross-process guard as the metrics write: a lifecycle transition must win. if ( @@ -4958,8 +4979,7 @@ export class C2DEngineDocker extends C2DEngine { */ private async sampleModelDownload( job: ServiceJob, - engine: ServiceEngineProfile, - startedAt: number + engine: ServiceEngineProfile ): Promise { try { const container = this.docker.getContainer(job.containerId) @@ -4969,20 +4989,12 @@ export class C2DEngineDocker extends C2DEngine { this.getContainerDiskUsage(job.containerId, path) ) if (fromManifest) return fromManifest - if (!engine.modelCachePath && !engine.comfyModelList) return undefined + if (!engine.modelCachePath) return undefined // Downloaded: what remains is the engine loading weights, which the cache says nothing about. // Returning nothing keeps the stored 100% record as it is. if (isModelDownloadComplete(job.modelDownload)) { return undefined } - if (engine.comfyModelList) { - const download = await sampleComfyModelDownload( - container, - !!job.outputBucketId, - startedAt - ) - return download ?? undefined - } const downloaded = await this.modelDownloadSampler.sample(job.serviceId, container) if (!downloaded) return undefined // Only a Hugging Face repo has a size the node can look up; a local path or an object-store diff --git a/src/components/c2d/downloadManifest.ts b/src/components/c2d/downloadManifest.ts index 62ca62113..383616302 100644 --- a/src/components/c2d/downloadManifest.ts +++ b/src/components/c2d/downloadManifest.ts @@ -20,7 +20,7 @@ import { buildModelDownload, fileSize } from './modelDownload.js' * manifest is untrusted: it only ever decides what this service's own progress report says. */ export const MANIFEST_PATH = '/tmp/ocean/downloads.tsv' -export const MAX_DOWNLOAD_ENTRIES = 64 +const MAX_DOWNLOAD_ENTRIES = 64 // How long after its container starts a service the node cannot probe is checked for a manifest. // A script writes it before its first download; setup steps before that can take minutes. export const MANIFEST_DISCOVERY_MS = 30 * 60 * 1000 @@ -37,15 +37,15 @@ export interface DownloadEntry { export type DirectoryBytes = (path: string) => Promise /** One small file from the container, via Docker's archive endpoint. Null when absent or too big. */ -export async function readContainerFile( +async function readContainerFile( container: Dockerode.Container, path: string -): Promise<{ text: string; mtime: number } | null> { +): Promise { try { const extract = tarStream.extract() pipeline(await container.getArchive({ path }), extract, () => {}) for await (const entry of extract) { - const { type, size, mtime } = entry.header + const { type, size } = entry.header if (type !== 'file' || (size ?? 0) > MAX_READ_BYTES) { return null } @@ -53,10 +53,7 @@ export async function readContainerFile( for await (const chunk of entry) { chunks.push(chunk) } - return { - text: Buffer.concat(chunks).toString('utf8'), - mtime: mtime?.getTime() ?? 0 - } + return Buffer.concat(chunks).toString('utf8') } return null } catch { @@ -152,7 +149,7 @@ export async function sampleDownloadManifest( if (!manifest) { return null } - const { entries, truncated } = parseDownloadManifest(manifest.text) + const { entries, truncated } = parseDownloadManifest(manifest) if (entries.length === 0) { return null } diff --git a/src/components/c2d/modelDownload.ts b/src/components/c2d/modelDownload.ts index e368a84fe..2a1dd1d5f 100644 --- a/src/components/c2d/modelDownload.ts +++ b/src/components/c2d/modelDownload.ts @@ -328,7 +328,7 @@ const BYTES_PER_PARAM: Record = { } const HF_MODEL_API = 'https://huggingface.co/api/models' -export const HF_TIMEOUT_MS = 8000 +const HF_TIMEOUT_MS = 8000 // One lookup per model for the life of the process: the answer cannot change for a given repo, and // this is read on the metrics cadence for every starting service. const totalBytesCache = new Map() @@ -362,7 +362,7 @@ export async function fetchModelTotalBytes( * Runs a Hub size lookup at most once per key: a size or a definitive null is kept for the life of * the process, while undefined (the Hub could not be asked) is retried after HUB_RETRY_AFTER_MS. */ -export async function cachedHubLookup( +async function cachedHubLookup( key: string, lookup: () => Promise ): Promise { diff --git a/src/components/c2d/serviceEngines.ts b/src/components/c2d/serviceEngines.ts index f752cf68d..21013eb87 100644 --- a/src/components/c2d/serviceEngines.ts +++ b/src/components/c2d/serviceEngines.ts @@ -45,8 +45,6 @@ export interface ServiceEngineProfile { * for a GGUF repo carrying a dozen quantizations differs by an order of magnitude. */ modelQuantFromCommand?: (cmd: string[] | undefined) => string | null - /** Progress is read from the model list a ComfyUI bundle script writes (see comfyDownload). */ - comfyModelList?: boolean } /** @@ -157,8 +155,7 @@ export const SERVICE_ENGINE_PROFILES: ServiceEngineProfile[] = [ matchesImage: (image) => /(^|\/)yanwk\/comfyui-boot$/.test(image), // ComfyUI binds its port only after the bundle script has fetched its models and every custom // node has loaded, so the first 200 means the UI is usable. - probe: { path: '/system_stats', port: 8188, expectStatus: [200] }, - comfyModelList: true + probe: { path: '/system_stats', port: 8188, expectStatus: [200] } } ] diff --git a/src/components/c2d/serviceReadiness.ts b/src/components/c2d/serviceReadiness.ts index ac370e57b..8c082cc3a 100644 --- a/src/components/c2d/serviceReadiness.ts +++ b/src/components/c2d/serviceReadiness.ts @@ -111,6 +111,9 @@ export const PROBE_PERIOD_SECONDS = 5 // Once ready the check keeps running — to catch an engine that dies without its container exiting — // but far more slowly: it is a liveness check at that point, not a wait. export const READY_PROBE_PERIOD_SECONDS = 30 +// A service still not ready this long after its container started is logged as a warning, once per +// container: usually a node that cannot reach it, sometimes just a very large download. +export const READINESS_WARN_AFTER_SECONDS = 15 * 60 const PROBE_TIMEOUT_MS = 2000 export interface ReadinessProbeResult { diff --git a/src/test/unit/service/comfyDownload.test.ts b/src/test/unit/service/comfyDownload.test.ts deleted file mode 100644 index 8c7983884..000000000 --- a/src/test/unit/service/comfyDownload.test.ts +++ /dev/null @@ -1,190 +0,0 @@ -import { expect } from 'chai' -import * as tarStream from 'tar-stream' -import type { ServiceJob } from '../../../@types/C2D/ServiceOnDemand.js' -import { resolveServiceEngine } from '../../../components/c2d/serviceEngines.js' -import { isModelDownloadComplete } from '../../../components/c2d/modelDownload.js' -import { - parseComfyModelList, - sampleComfyModelDownload -} from '../../../components/c2d/comfyDownload.js' - -const HF = 'https://huggingface.co' - -describe('ComfyUI engine profile', () => { - const job = (image: string) => ({ image, exposedPorts: [8188] }) as ServiceJob - - it('recognizes the comfyui-boot image under any tag', () => { - expect(resolveServiceEngine(job('yanwk/comfyui-boot'))?.id).to.equal('comfyui') - expect( - resolveServiceEngine(job('yanwk/comfyui-boot:cu130-megapak-pt211'))?.id - ).to.equal('comfyui') - expect(resolveServiceEngine(job('someone/comfyui-boot'))).to.equal(null) - }) -}) - -describe('parseComfyModelList', () => { - it('keeps plain Hugging Face file URLs and derives the script destination', () => { - const text = [ - `vae\t${HF}/org/repo/resolve/main/split_files/vae/a.safetensors`, - `loras\t${HF}/org/repo/resolve/main/b.safetensors`, - `loras\t${HF}/org/repo/resolve/main/b.safetensors`, - '' - ].join('\n') - expect(parseComfyModelList(text, '/tmp/comfy')).to.deep.equal({ - entries: [ - { - url: `${HF}/org/repo/resolve/main/split_files/vae/a.safetensors`, - path: '/tmp/comfy/models/vae/a.safetensors' - }, - { - url: `${HF}/org/repo/resolve/main/b.safetensors`, - path: '/tmp/comfy/models/loras/b.safetensors' - } - ], - truncated: false - }) - }) - - it('caps the list and flags one that names more files than the cap', () => { - const line = (i: number) => `vae\t${HF}/org/repo/resolve/main/f${i}.safetensors` - const capped = Array.from({ length: 64 }, (_, i) => line(i)) - const exact = parseComfyModelList([...capped, line(0)].join('\n'), '/tmp/comfy') - expect(exact.entries).to.have.length(64) - expect(exact.truncated).to.equal(false) - - const over = parseComfyModelList([...capped, line(64)].join('\n'), '/tmp/comfy') - expect(over.entries).to.have.length(64) - expect(over.truncated).to.equal(true) - }) - - it('drops anything that is not a plain Hugging Face file URL or a safe directory', () => { - const text = [ - `../etc\t${HF}/org/repo/resolve/main/a.safetensors`, - `vae\thttps://example.com/org/repo/resolve/main/a.safetensors`, - `vae\t${HF}/org/repo/resolve/main/a.safetensors?download=true`, - `vae\t${HF}:8443/org/repo/resolve/main/a.safetensors`, - `vae\t${HF}/org/repo/resolve/main/../a.safetensors`, - `vae\t${HF}/org/repo/blob/main/a.safetensors`, - `vae\t${HF}/org/repo/resolve/main/a.bin`, - 'vae' - ].join('\n') - expect(parseComfyModelList(text, '/tmp/comfy').entries).to.deep.equal([]) - }) -}) - -describe('sampleComfyModelDownload', () => { - const realFetch = globalThis.fetch - const startedAt = Date.parse('2026-10-07T10:00:00Z') - - afterEach(() => { - globalThis.fetch = realFetch - }) - - // A container double: `getArchive()` returns the list as a tar, `infoArchive()` answers stats. - function fakeContainer(list: string, mtime: Date, files: Record) { - return { - getArchive: () => { - const pack = tarStream.pack() - pack.entry({ name: '.models.tsv', mtime }, list) - pack.finalize() - return Promise.resolve(pack) - }, - infoArchive: ({ path }: { path: string }) => { - if (files[path] === undefined) { - return Promise.reject(new Error('not found')) - } - const header = Buffer.from(JSON.stringify({ size: files[path] })).toString( - 'base64' - ) - return Promise.resolve({ headers: { 'x-docker-container-path-stat': header } }) - } - } as any - } - - function hubSizes(sizes: Record) { - globalThis.fetch = ((url: string) => - Promise.resolve( - new Response(null, { - status: 302, - headers: { 'x-linked-size': String(sizes[url]) } - }) - )) as typeof fetch - } - - it('measures finished and partial files against the listed total', async () => { - const a = `${HF}/sample-test/repo/resolve/main/a.safetensors` - const b = `${HF}/sample-test/repo/resolve/main/b.safetensors` - const c = `${HF}/sample-test/repo/resolve/main/c.safetensors` - hubSizes({ [a]: 1000, [b]: 2000, [c]: 1000 }) - const container = fakeContainer( - `vae\t${a}\nloras\t${b}\nloras\t${c}\n`, - new Date(startedAt + 5000), - { - '/data/outputs/comfy/models/vae/a.safetensors': 1000, - '/data/outputs/comfy/models/loras/b.safetensors.part': 500 - } - ) - const download = await sampleComfyModelDownload(container, true, startedAt) - expect(download).to.include({ - downloadedBytes: 1500, - totalBytes: 4000, - percent: 38, - filesComplete: 1, - filesInFlight: 1, - filesTotal: 3 - }) - expect(isModelDownloadComplete(download)).to.equal(false) - }) - - it('is complete once every listed file is finished', async () => { - const a = `${HF}/complete-test/repo/resolve/main/a.safetensors` - hubSizes({ [a]: 1000 }) - const container = fakeContainer(`vae\t${a}\n`, new Date(startedAt), { - '/tmp/comfy/models/vae/a.safetensors': 1000 - }) - const download = await sampleComfyModelDownload(container, false, startedAt) - expect(isModelDownloadComplete(download)).to.equal(true) - }) - - it('ignores a list left in the bucket by an earlier launch', async () => { - const a = `${HF}/stale-test/repo/resolve/main/a.safetensors` - const container = fakeContainer(`vae\t${a}\n`, new Date(startedAt - 60_000), {}) - expect(await sampleComfyModelDownload(container, true, startedAt)).to.equal(null) - }) - - it('reports no total when a size is unknown, and nothing when there is no list', async () => { - const a = `${HF}/unknown-test/repo/resolve/main/a.safetensors` - globalThis.fetch = (() => - Promise.resolve(new Response(null, { status: 401 }))) as typeof fetch - const container = fakeContainer(`vae\t${a}\n`, new Date(startedAt), {}) - const download = await sampleComfyModelDownload(container, false, startedAt) - expect(download?.totalBytes).to.equal(undefined) - expect(download?.filesTotal).to.equal(1) - - const empty = { getArchive: () => Promise.reject(new Error('no such file')) } as any - expect(await sampleComfyModelDownload(empty, false, startedAt)).to.equal(null) - }) - - it('reports bytes only for a list longer than the cap, so it never completes early', async () => { - const url = (i: number) => `${HF}/overflow-test/repo/resolve/main/f${i}.safetensors` - const list = Array.from({ length: 65 }, (_, i) => `vae\t${url(i)}`).join('\n') - const files = Object.fromEntries( - Array.from({ length: 64 }, (_, i) => [ - `/tmp/comfy/models/vae/f${i}.safetensors`, - 10 - ]) - ) - let lookups = 0 - globalThis.fetch = (() => { - lookups++ - return Promise.resolve(new Response(null, { status: 302 })) - }) as typeof fetch - const container = fakeContainer(list, new Date(startedAt), files) - const download = await sampleComfyModelDownload(container, false, startedAt) - expect(download).to.include({ downloadedBytes: 640, filesComplete: 64 }) - expect(download?.filesTotal).to.equal(undefined) - expect(download?.totalBytes).to.equal(undefined) - expect(isModelDownloadComplete(download)).to.equal(false) - expect(lookups).to.equal(0) - }) -}) diff --git a/src/test/unit/service/readinessWarning.test.ts b/src/test/unit/service/readinessWarning.test.ts new file mode 100644 index 000000000..c6641b0ad --- /dev/null +++ b/src/test/unit/service/readinessWarning.test.ts @@ -0,0 +1,73 @@ +import { expect } from 'chai' +import sinon from 'sinon' +import { C2DEngineDocker } from '../../../components/c2d/compute_engine_docker.js' +import { ServiceStatusNumber } from '../../../@types/C2D/ServiceOnDemand.js' +import { CORE_LOGGER } from '../../../utils/logging/common.js' + +// A ComfyUI service whose port the node cannot reach (port 1 refuses at once), so every probe +// reads "waiting" — the case the warning exists for. +describe('warning for a service that stays not ready', () => { + let warn: sinon.SinonStub + beforeEach(() => { + warn = sinon.stub(CORE_LOGGER, 'warn') + }) + afterEach(() => { + warn.restore() + }) + + function engine() { + const e: any = Object.create(C2DEngineDocker.prototype) + e.serviceOpsInFlight = new Set() + e.serviceProbeUrls = new Map() + e.manifestCheckedAt = new Map() + e.readinessWarned = new Map() + e.docker = { + getContainer: () => ({ + getArchive: () => Promise.reject(new Error('no such file')) + }) + } + e.db = { + isServiceLocked: sinon.stub().resolves(false), + updateServiceJobReadiness: sinon.stub().resolves(true) + } + return e + } + + const job = (containerId = 'container-1') => ({ + serviceId: 'svc-1', + owner: '0x0000000000000000000000000000000000000001', + clusterHash: '0xcluster', + containerId, + image: 'yanwk/comfyui-boot', + exposedPorts: [8188], + endpoints: [{ containerPort: 8188, hostPort: 1, url: 'http://127.0.0.1:1' }], + status: ServiceStatusNumber.Running + }) + const startedAgo = (minutes: number) => ({ + State: { StartedAt: new Date(Date.now() - minutes * 60_000).toISOString() }, + NetworkSettings: { Networks: {} } + }) + const readinessWarnings = () => + warn.getCalls().filter((call) => String(call.args[0]).startsWith('[readiness]')) + + it('warns once per container after 15 minutes', async () => { + const e = engine() + await e.probeServiceReadiness(job(), startedAgo(16)) + await e.probeServiceReadiness(job(), startedAgo(16)) + expect(readinessWarnings()).to.have.length(1) + expect(readinessWarnings()[0].args[0]).to.include('svc-1 (comfyui) not ready 16 min') + }) + + it('stays quiet before then', async () => { + const e = engine() + await e.probeServiceReadiness(job(), startedAgo(5)) + expect(readinessWarnings()).to.have.length(0) + }) + + it('warns again for the new container after a restart', async () => { + const e = engine() + await e.probeServiceReadiness(job('container-1'), startedAgo(16)) + await e.probeServiceReadiness(job('container-2'), startedAgo(16)) + expect(readinessWarnings()).to.have.length(2) + }) +}) diff --git a/src/test/unit/service/serviceReadiness.test.ts b/src/test/unit/service/serviceReadiness.test.ts index 7398f7a4a..92fd1a35c 100644 --- a/src/test/unit/service/serviceReadiness.test.ts +++ b/src/test/unit/service/serviceReadiness.test.ts @@ -180,6 +180,18 @@ describe('resolveServiceEngine', () => { }) }) +describe('resolveServiceEngine — ComfyUI', () => { + const job = (image: string) => ({ image, exposedPorts: [8188] }) as ServiceJob + + it('recognizes the comfyui-boot image under any tag', () => { + expect(resolveServiceEngine(job('yanwk/comfyui-boot'))?.id).to.equal('comfyui') + expect( + resolveServiceEngine(job('yanwk/comfyui-boot:cu130-megapak-pt211'))?.id + ).to.equal('comfyui') + expect(resolveServiceEngine(job('someone/comfyui-boot'))).to.equal(null) + }) +}) + describe('resolveServiceEngine — llama.cpp', () => { const job = (image: string, cmd?: string[]) => ({ image, dockerCmd: cmd, exposedPorts: [8000, 8080] }) as ServiceJob From 2acf016050e543a1ecf7608457f74fec4a826074 Mon Sep 17 00:00:00 2001 From: andreip135 <338078672+andreip135@users.noreply.github.com> Date: Fri, 9 Oct 2026 01:38:29 +0300 Subject: [PATCH 6/8] measure manifest folders at most every 15s --- src/components/c2d/compute_engine_docker.ts | 35 +++++++++++++++++-- .../service/unprobedServiceDownload.test.ts | 10 ++++++ 2 files changed, 42 insertions(+), 3 deletions(-) diff --git a/src/components/c2d/compute_engine_docker.ts b/src/components/c2d/compute_engine_docker.ts index 50601f4bf..7b22f75d6 100755 --- a/src/components/c2d/compute_engine_docker.ts +++ b/src/components/c2d/compute_engine_docker.ts @@ -150,6 +150,9 @@ const SERVICE_LOCK_STALE_MS = 120_000 // A `du` inside a container that never finishes (a huge or slow folder, or a container that replaced // the binary) must not hold the loop that measures it, nor the node's shutdown that drains it. const DISK_USAGE_TIMEOUT_MS = 15_000 +// A manifest folder is measured with `du` at most this often, while the manifest itself is re-read +// on every sample to catch appended entries. +const MANIFEST_DIR_MEASURE_MS = 15_000 // Identifies one engine instance as a service-lock holder across processes. pid alone is // not enough: pids are reused, and one process can host several engine instances. @@ -224,6 +227,11 @@ export class C2DEngineDocker extends C2DEngine { private manifestCheckedAt: Map = new Map() // serviceId -> the container a "still not ready" warning was logged for (once per container). private readinessWarned: Map = new Map() + // serviceId -> the last `du` of each manifest folder in its current container. + private manifestDirSizes: Map< + string, + { containerId: string; sizes: Map } + > = new Map() // serviceId -> its readiness probe (+ model-download sample) still running in the background. // The probe is launched fire-and-forget so a slow engine, Docker daemon or Hub never holds up // InternalLoop; this keeps one probe per service at a time and lets stop() drain them. @@ -4675,7 +4683,8 @@ export class C2DEngineDocker extends C2DEngine { for (const state of [ this.serviceProbeUrls, this.manifestCheckedAt, - this.readinessWarned + this.readinessWarned, + this.manifestDirSizes ]) { for (const serviceId of state.keys()) { if (!running.has(serviceId)) { @@ -4940,7 +4949,7 @@ export class C2DEngineDocker extends C2DEngine { this.manifestCheckedAt.set(job.serviceId, now) const modelDownload = await sampleDownloadManifest( this.docker.getContainer(job.containerId), - (path) => this.getContainerDiskUsage(job.containerId, path) + (path) => this.manifestDirectoryBytes(job, path) ) if (!modelDownload) { return @@ -4965,6 +4974,26 @@ export class C2DEngineDocker extends C2DEngine { ) } + /** A manifest folder's size, re-measured with `du` at most every MANIFEST_DIR_MEASURE_MS. */ + private async manifestDirectoryBytes( + job: ServiceJob, + path: string + ): Promise { + let cache = this.manifestDirSizes.get(job.serviceId) + if (!cache || cache.containerId !== job.containerId) { + cache = { containerId: job.containerId, sizes: new Map() } + this.manifestDirSizes.set(job.serviceId, cache) + } + const now = Date.now() + const last = cache.sizes.get(path) + if (last && now - last.at < MANIFEST_DIR_MEASURE_MS) { + return last.bytes + } + const bytes = await this.getContainerDiskUsage(job.containerId, path) + cache.sizes.set(path, { bytes, at: now }) + return bytes + } + /** * How much of its model the container has pulled down: from the launch script's download * manifest when it writes one, otherwise from the engine's own cache. @@ -4986,7 +5015,7 @@ export class C2DEngineDocker extends C2DEngine { // The launch script's own manifest, when it writes one (see downloadManifest). Read before // the completion check below: a script can append to it after its first downloads finish. const fromManifest = await sampleDownloadManifest(container, (path) => - this.getContainerDiskUsage(job.containerId, path) + this.manifestDirectoryBytes(job, path) ) if (fromManifest) return fromManifest if (!engine.modelCachePath) return undefined diff --git a/src/test/unit/service/unprobedServiceDownload.test.ts b/src/test/unit/service/unprobedServiceDownload.test.ts index a6e7b57dd..42a1d2044 100644 --- a/src/test/unit/service/unprobedServiceDownload.test.ts +++ b/src/test/unit/service/unprobedServiceDownload.test.ts @@ -14,6 +14,7 @@ describe('download progress for a service the node cannot probe', () => { const engine: any = Object.create(C2DEngineDocker.prototype) engine.serviceOpsInFlight = new Set() engine.manifestCheckedAt = new Map() + engine.manifestDirSizes = new Map() engine.getArchive = sinon.stub().callsFake(() => { if (manifest === null) { return Promise.reject(new Error('no such file')) @@ -96,6 +97,15 @@ describe('download progress for a service the node cannot probe', () => { sinon.assert.notCalled(engine.getArchive) }) + it('measures a manifest folder at most once per period, and again in a new container', async () => { + const engine = engineWith(null, 400) + await engine.manifestDirectoryBytes(job(), CACHE) + await engine.manifestDirectoryBytes(job(), CACHE) + sinon.assert.calledOnce(engine.getContainerDiskUsage) + await engine.manifestDirectoryBytes(job({ containerId: 'container-2' }), CACHE) + sinon.assert.calledTwice(engine.getContainerDiskUsage) + }) + it('writes nothing for a service that has no manifest', async () => { const engine = engineWith(null, null) await engine.probeServiceReadiness(job(), startedAgo(60_000)) From 5b2d88514bd6f3ceeecf7efa0da4dcba2e921592 Mon Sep 17 00:00:00 2001 From: andreip135 <338078672+andreip135@users.noreply.github.com> Date: Fri, 9 Oct 2026 02:16:29 +0300 Subject: [PATCH 7/8] find du size after error lines; time out only manifest du; cap manifest folders and sampling time --- src/@types/C2D/ServiceOnDemand.ts | 2 +- src/components/c2d/compute_engine_docker.ts | 47 +++++++++++++------ src/components/c2d/downloadManifest.ts | 34 +++++++++----- src/components/c2d/modelDownload.ts | 15 +----- src/components/database/sqliteCompute.ts | 4 +- .../unit/service/downloadManifest.test.ts | 18 +++++++ .../service/unprobedServiceDownload.test.ts | 27 ++++++++++- 7 files changed, 105 insertions(+), 42 deletions(-) diff --git a/src/@types/C2D/ServiceOnDemand.ts b/src/@types/C2D/ServiceOnDemand.ts index 6ee02c850..ce0c36000 100644 --- a/src/@types/C2D/ServiceOnDemand.ts +++ b/src/@types/C2D/ServiceOnDemand.ts @@ -193,7 +193,7 @@ export interface ServiceModelDownload { percent?: number filesComplete: number filesInFlight: number // the hub fetches files in parallel, so only the aggregate is meaningful - filesTotal?: number // files the service will download, when it lists them up front (ComfyUI) + filesTotal?: number // items the service will download, when its download manifest lists them updatedAt: number // Unix ms } diff --git a/src/components/c2d/compute_engine_docker.ts b/src/components/c2d/compute_engine_docker.ts index 7b22f75d6..722222ae0 100755 --- a/src/components/c2d/compute_engine_docker.ts +++ b/src/components/c2d/compute_engine_docker.ts @@ -100,7 +100,11 @@ import { isModelDownloadComplete, ModelDownloadSampler } from './modelDownload.js' -import { MANIFEST_DISCOVERY_MS, sampleDownloadManifest } from './downloadManifest.js' +import { + MANIFEST_DISCOVERY_MS, + MANIFEST_SAMPLING_MS, + sampleDownloadManifest +} from './downloadManifest.js' import type { DockerMountObject } from '../../@types/PersistentStorage.js' import { resolveServiceImage } from './serviceResourceMatching.js' import { @@ -147,9 +151,10 @@ const trivyImage = 'aquasec/trivy:0.69.3' // Use pinned versions for safety // doesn't get a live operation's resources torn down under it. const SERVICE_LOCK_HEARTBEAT_MS = 30_000 const SERVICE_LOCK_STALE_MS = 120_000 -// A `du` inside a container that never finishes (a huge or slow folder, or a container that replaced -// the binary) must not hold the loop that measures it, nor the node's shutdown that drains it. -const DISK_USAGE_TIMEOUT_MS = 15_000 +// How long the node waits for `du` on a service's manifest folder. The container is user-controlled +// and the measurement is only progress reporting, so it must not hold the probe or the node's +// shutdown that drains it. Compute disk-quota checks keep waiting for the answer. +const MANIFEST_DU_TIMEOUT_MS = 15_000 // A manifest folder is measured with `du` at most this often, while the manifest itself is re-read // on every sample to catch appended entries. const MANIFEST_DIR_MEASURE_MS = 15_000 @@ -3083,7 +3088,8 @@ export class C2DEngineDocker extends C2DEngine { // never report an unmeasurable container as "0 bytes used". private async getContainerDiskUsage( containerName: string, - path: string = '/data' + path: string = '/data', + timeoutMs?: number ): Promise { try { const container = this.docker.getContainer(containerName) @@ -3117,13 +3123,15 @@ export class C2DEngineDocker extends C2DEngine { read.catch(() => {}) let timer: NodeJS.Timeout const timedOut = new Promise((resolve) => { - timer = setTimeout(() => resolve(false), DISK_USAGE_TIMEOUT_MS) + if (timeoutMs !== undefined) { + timer = setTimeout(() => resolve(false), timeoutMs) + } }) try { if (!(await Promise.race([read, timedOut]))) { stream.destroy() CORE_LOGGER.warn( - `du ${path} in ${containerName} did not finish within ${DISK_USAGE_TIMEOUT_MS / 1000}s` + `du ${path} in ${containerName} did not finish within ${timeoutMs / 1000}s` ) return null } @@ -3133,8 +3141,9 @@ export class C2DEngineDocker extends C2DEngine { const output = Buffer.concat(chunks).toString() - // `du -sb` prints "\t"; anything else (an error) is not a size. - const match = output.match(/^\s*(\d+)\s/) + // `du -sb` prints "\t", after any "du: cannot access" lines (e.g. for /proc), + // which the TTY merges into the same output. + const match = output.match(/^(\d+)\t/m) return match ? parseInt(match[1], 10) : null } catch (error) { CORE_LOGGER.error( @@ -4918,7 +4927,8 @@ export class C2DEngineDocker extends C2DEngine { * manifest says, with no readiness, so Running stays the signal that the service is usable. * * Throttled like the probe. A service with no manifest MANIFEST_DISCOVERY_MS after its container - * started is not checked again, and sampling stops once the listed download is complete — with no + * started is not checked again, and sampling stops once the listed download is complete or after + * MANIFEST_SAMPLING_MS, whichever comes first — with no * readiness to end it, it would otherwise run for the whole session. */ private async sampleUnprobedServiceDownload( @@ -4934,9 +4944,8 @@ export class C2DEngineDocker extends C2DEngine { const now = Date.now() const startedAt = Date.parse(details.State?.StartedAt ?? '') if ( - !job.modelDownload && Number.isFinite(startedAt) && - now - startedAt > MANIFEST_DISCOVERY_MS + now - startedAt > (job.modelDownload ? MANIFEST_SAMPLING_MS : MANIFEST_DISCOVERY_MS) ) { return } @@ -4989,7 +4998,11 @@ export class C2DEngineDocker extends C2DEngine { if (last && now - last.at < MANIFEST_DIR_MEASURE_MS) { return last.bytes } - const bytes = await this.getContainerDiskUsage(job.containerId, path) + const bytes = await this.getContainerDiskUsage( + job.containerId, + path, + MANIFEST_DU_TIMEOUT_MS + ) cache.sizes.set(path, { bytes, at: now }) return bytes } @@ -5017,8 +5030,12 @@ export class C2DEngineDocker extends C2DEngine { const fromManifest = await sampleDownloadManifest(container, (path) => this.manifestDirectoryBytes(job, path) ) - if (fromManifest) return fromManifest - if (!engine.modelCachePath) return undefined + if (fromManifest) { + return fromManifest + } + if (!engine.modelCachePath) { + return undefined + } // Downloaded: what remains is the engine loading weights, which the cache says nothing about. // Returning nothing keeps the stored 100% record as it is. if (isModelDownloadComplete(job.modelDownload)) { diff --git a/src/components/c2d/downloadManifest.ts b/src/components/c2d/downloadManifest.ts index 383616302..555f51159 100644 --- a/src/components/c2d/downloadManifest.ts +++ b/src/components/c2d/downloadManifest.ts @@ -21,11 +21,16 @@ import { buildModelDownload, fileSize } from './modelDownload.js' */ export const MANIFEST_PATH = '/tmp/ocean/downloads.tsv' const MAX_DOWNLOAD_ENTRIES = 64 +// Each `dir` entry costs a `du` inside the container, so far fewer of them are measured. +const MAX_DIR_ENTRIES = 8 // How long after its container starts a service the node cannot probe is checked for a manifest. // A script writes it before its first download; setup steps before that can take minutes. export const MANIFEST_DISCOVERY_MS = 30 * 60 * 1000 +// The longest a service the node cannot probe is sampled after its container starts, so a download +// that never reads as complete (an unknown size, a file that never arrives) is not sampled forever. +export const MANIFEST_SAMPLING_MS = 6 * 60 * 60 * 1000 const MAX_READ_BYTES = 64 * 1024 -const BYTES = /^\d{1,15}$/ +const BYTES = /^\d{1,13}$/ export interface DownloadEntry { kind: 'file' | 'dir' @@ -62,14 +67,15 @@ async function readContainerFile( } /** - * The manifest's entries, capped at MAX_DOWNLOAD_ENTRIES. `truncated` says it names more than - * that, so the entries are not the whole download. + * The manifest's entries, capped at MAX_DOWNLOAD_ENTRIES (MAX_DIR_ENTRIES of them `dir`). `truncated` + * says it names more than that, so the entries are not the whole download. */ export function parseDownloadManifest(text: string): { entries: DownloadEntry[] truncated: boolean } { const entries = new Map() + let dirs = 0 for (const line of text.split('\n')) { const [kind, bytes, path] = line.trim().split('\t') if ( @@ -81,9 +87,15 @@ export function parseDownloadManifest(text: string): { ) { continue } - if (entries.size >= MAX_DOWNLOAD_ENTRIES) { + if ( + entries.size >= MAX_DOWNLOAD_ENTRIES || + (kind === 'dir' && dirs >= MAX_DIR_ENTRIES) + ) { return { entries: [...entries.values()], truncated: true } } + if (kind === 'dir') { + dirs++ + } entries.set(path, { kind, path, bytes: Number(bytes) || null }) } return { entries: [...entries.values()], truncated: false } @@ -117,17 +129,17 @@ export async function measureDownloads( } continue } - // Partial first: a file renamed between the two checks is then still found finished. - const partial = await fileSize(container, `${path}.part`) - if (partial !== null) { - downloadedBytes += partial - inFlight++ - continue - } + // Finished first, so a stale .part left beside it cannot hold the file in flight forever. const finished = await fileSize(container, path) if (finished !== null) { downloadedBytes += finished files++ + continue + } + const partial = await fileSize(container, `${path}.part`) + if (partial !== null) { + downloadedBytes += partial + inFlight++ } } const totalBytes = diff --git a/src/components/c2d/modelDownload.ts b/src/components/c2d/modelDownload.ts index 2a1dd1d5f..d8b33fc56 100644 --- a/src/components/c2d/modelDownload.ts +++ b/src/components/c2d/modelDownload.ts @@ -355,23 +355,12 @@ export async function fetchModelTotalBytes( quant?: string ): Promise { const key = quant ? `${modelId}:${quant}` : modelId - return await cachedHubLookup(key, () => lookupModelTotalBytes(modelId, quant)) -} - -/** - * Runs a Hub size lookup at most once per key: a size or a definitive null is kept for the life of - * the process, while undefined (the Hub could not be asked) is retried after HUB_RETRY_AFTER_MS. - */ -async function cachedHubLookup( - key: string, - lookup: () => Promise -): Promise { if (totalBytesCache.has(key)) return totalBytesCache.get(key) ?? null const failedAt = hubFailedAt.get(key) if (failedAt !== undefined && Date.now() - failedAt < HUB_RETRY_AFTER_MS) { return null } - const total = await lookup() + const total = await lookupModelTotalBytes(modelId, quant) if (total === undefined) { hubFailedAt.set(key, Date.now()) return null @@ -482,7 +471,7 @@ async function fetchGgufFileBytes( * is an estimate (dtype × parameters), so on its own it can read 100% while the last shard is still * being written; no partial files left is what makes it final. A record with no total never * completes here, and sampling simply carries on until the service is ready. A record that knows - * its file count (a ComfyUI model list) is complete once every listed file is. + * its file count (a download manifest) is complete once every listed item is. */ export function isModelDownloadComplete( download: ServiceModelDownload | undefined diff --git a/src/components/database/sqliteCompute.ts b/src/components/database/sqliteCompute.ts index b034fb111..a801cf2f4 100644 --- a/src/components/database/sqliteCompute.ts +++ b/src/components/database/sqliteCompute.ts @@ -423,7 +423,9 @@ export class SQLiteCompute implements ComputeDatabaseProvider { modelDownload?: ServiceModelDownload ): Promise { return this.patchServiceJobBody(serviceId, expected, 'readiness', (body) => { - if (readiness) body.readiness = readiness + if (readiness) { + body.readiness = readiness + } // Only overwritten when a fresh sample was taken: once the engine is ready the walk stops, // and the last figures stay as the record of what was downloaded. if (modelDownload) body.modelDownload = modelDownload diff --git a/src/test/unit/service/downloadManifest.test.ts b/src/test/unit/service/downloadManifest.test.ts index 2acbb1b7d..7372252ba 100644 --- a/src/test/unit/service/downloadManifest.test.ts +++ b/src/test/unit/service/downloadManifest.test.ts @@ -44,6 +44,15 @@ describe('parseDownloadManifest', () => { ]) }) + it('measures at most 8 folders and flags a manifest that names more', () => { + const dir = (i: number) => `dir\t1\t/data/hub/models--org--m${i}` + const eight = Array.from({ length: 8 }, (_, i) => dir(i)) + expect(parseDownloadManifest(eight.join('\n')).truncated).to.equal(false) + const over = parseDownloadManifest([...eight, dir(8)].join('\n')) + expect(over.entries).to.have.length(8) + expect(over.truncated).to.equal(true) + }) + it('caps the entries and flags a manifest that names more', () => { const line = (i: number) => `file\t1\t${MODELS}/f${i}` const capped = Array.from({ length: 64 }, (_, i) => line(i)) @@ -119,6 +128,15 @@ describe('sampleDownloadManifest', () => { expect(isModelDownloadComplete(download)).to.equal(true) }) + it('counts a finished file as done even with a stale .part beside it', async () => { + const container = fakeContainer(`file\t1000\t${MODELS}/vae/a.safetensors`, { + [`${MODELS}/vae/a.safetensors`]: 1000, + [`${MODELS}/vae/a.safetensors.part`]: 10 + }) + const download = await sampleDownloadManifest(container, () => Promise.resolve(null)) + expect(download).to.include({ filesComplete: 1, filesInFlight: 0, percent: 100 }) + }) + it('reports no total when a size is unknown', async () => { const container = fakeContainer(`file\t0\t${MODELS}/vae/a.safetensors`, {}) const download = await sampleDownloadManifest(container, () => Promise.resolve(null)) diff --git a/src/test/unit/service/unprobedServiceDownload.test.ts b/src/test/unit/service/unprobedServiceDownload.test.ts index 42a1d2044..67bd9e3db 100644 --- a/src/test/unit/service/unprobedServiceDownload.test.ts +++ b/src/test/unit/service/unprobedServiceDownload.test.ts @@ -106,6 +106,22 @@ describe('download progress for a service the node cannot probe', () => { sinon.assert.calledTwice(engine.getContainerDiskUsage) }) + it('stops sampling six hours after the container started, complete or not', async () => { + const engine = engineWith(`dir\t0\t${CACHE}\n`, 400) + const unfinished = { + downloadedBytes: 400, + filesComplete: 0, + filesInFlight: 1, + filesTotal: 1, + updatedAt: Date.now() + } + await engine.probeServiceReadiness( + job({ modelDownload: unfinished }), + startedAgo(6 * 60 * 60 * 1000 + 1) + ) + sinon.assert.notCalled(engine.getArchive) + }) + it('writes nothing for a service that has no manifest', async () => { const engine = engineWith(null, null) await engine.probeServiceReadiness(job(), startedAgo(60_000)) @@ -133,6 +149,15 @@ describe('getContainerDiskUsage', () => { expect(await engine.getContainerDiskUsage('container-1', path)).to.equal(5004096) }) + it('finds the size after the error lines du prints for /proc', async () => { + const engine = engineWithOutput( + "du: cannot access '/proc/7/task/7/fd/3': No such file or directory\r\n" + + "du: cannot access '/proc/7/fd/4': No such file or directory\r\n" + + '204542432\t/\r\n' + ) + expect(await engine.getContainerDiskUsage('container-1', '/')).to.equal(204542432) + }) + it('never reads an error message as a size', async () => { const engine = engineWithOutput( "du: cannot access '/data/models--org--model-2 1': No such file or directory\r\n" @@ -159,7 +184,7 @@ describe('getContainerDiskUsage timeout', () => { exec: sinon.stub().resolves({ start: sinon.stub().resolves(stream) }) }) } - const result = engine.getContainerDiskUsage('container-1', '/data/hung') + const result = engine.getContainerDiskUsage('container-1', '/data/hung', 15_000) await clock.tickAsync(15_000) expect(await result).to.equal(null) expect(stream.destroyed).to.equal(true) From 7c4842c044d820ac8568e5e39ac3b55ab596ddd4 Mon Sep 17 00:00:00 2001 From: andreip135 <338078672+andreip135@users.noreply.github.com> Date: Thu, 15 Oct 2026 15:52:38 +0300 Subject: [PATCH 8/8] cleanup --- src/@types/C2D/ServiceOnDemand.ts | 2 +- src/components/c2d/compute_engine_docker.ts | 156 +++++++++----------- src/components/c2d/downloadManifest.ts | 79 ++++------ src/components/c2d/modelDownload.ts | 4 +- src/components/c2d/serviceEngines.ts | 9 +- src/components/c2d/serviceReadiness.ts | 3 +- src/components/database/sqliteCompute.ts | 2 +- 7 files changed, 101 insertions(+), 154 deletions(-) diff --git a/src/@types/C2D/ServiceOnDemand.ts b/src/@types/C2D/ServiceOnDemand.ts index ce0c36000..a038c4e31 100644 --- a/src/@types/C2D/ServiceOnDemand.ts +++ b/src/@types/C2D/ServiceOnDemand.ts @@ -193,7 +193,7 @@ export interface ServiceModelDownload { percent?: number filesComplete: number filesInFlight: number // the hub fetches files in parallel, so only the aggregate is meaningful - filesTotal?: number // items the service will download, when its download manifest lists them + filesTotal?: number // from the download manifest updatedAt: number // Unix ms } diff --git a/src/components/c2d/compute_engine_docker.ts b/src/components/c2d/compute_engine_docker.ts index 722222ae0..b3612c398 100755 --- a/src/components/c2d/compute_engine_docker.ts +++ b/src/components/c2d/compute_engine_docker.ts @@ -151,12 +151,8 @@ const trivyImage = 'aquasec/trivy:0.69.3' // Use pinned versions for safety // doesn't get a live operation's resources torn down under it. const SERVICE_LOCK_HEARTBEAT_MS = 30_000 const SERVICE_LOCK_STALE_MS = 120_000 -// How long the node waits for `du` on a service's manifest folder. The container is user-controlled -// and the measurement is only progress reporting, so it must not hold the probe or the node's -// shutdown that drains it. Compute disk-quota checks keep waiting for the answer. +// `du` on a manifest folder runs in a user-controlled container, so it must not hold the probe. const MANIFEST_DU_TIMEOUT_MS = 15_000 -// A manifest folder is measured with `du` at most this often, while the manifest itself is re-read -// on every sample to catch appended entries. const MANIFEST_DIR_MEASURE_MS = 15_000 // Identifies one engine instance as a service-lock holder across processes. pid alone is @@ -228,14 +224,14 @@ export class C2DEngineDocker extends C2DEngine { private serviceProbeUrls: Map = new Map() // Per-service model-download sampling state (which cache files to stat, when to re-list them). private modelDownloadSampler = new ModelDownloadSampler() - // serviceId -> when a service the node cannot probe was last checked for a download manifest. + // serviceId -> last manifest check, for services the node cannot probe. private manifestCheckedAt: Map = new Map() - // serviceId -> the container a "still not ready" warning was logged for (once per container). + // serviceId -> containerId the "still not ready" warning was logged for. private readinessWarned: Map = new Map() - // serviceId -> the last `du` of each manifest folder in its current container. + // serviceId -> manifest folder path -> last `du`. private manifestDirSizes: Map< string, - { containerId: string; sizes: Map } + Map > = new Map() // serviceId -> its readiness probe (+ model-download sample) still running in the background. // The probe is launched fire-and-forget so a slow engine, Docker daemon or Hub never holds up @@ -3101,8 +3097,7 @@ export class C2DEngineDocker extends C2DEngine { return null } - // With a TTY the output arrives as plain text. Without one, Docker frames it with 8-byte - // headers whose length byte can read as an ASCII digit and get prepended to the size. + // Without a TTY, Docker's 8-byte stream headers can read as digits prepended to the size. const exec = await container.exec({ Cmd: ['du', '-sb', path], AttachStdout: true, @@ -3112,37 +3107,36 @@ export class C2DEngineDocker extends C2DEngine { const stream = await exec.start({ Detach: false, Tty: true }) + let timedOut = false + const timer = + timeoutMs !== undefined + ? setTimeout(() => { + timedOut = true + stream.destroy() + }, timeoutMs) + : undefined const chunks: Buffer[] = [] - const read = (async () => { + try { for await (const chunk of stream) { chunks.push(chunk as Buffer) } - return true - })() - // Destroying the stream on timeout rejects the read after the race has moved on. - read.catch(() => {}) - let timer: NodeJS.Timeout - const timedOut = new Promise((resolve) => { - if (timeoutMs !== undefined) { - timer = setTimeout(() => resolve(false), timeoutMs) - } - }) - try { - if (!(await Promise.race([read, timedOut]))) { - stream.destroy() - CORE_LOGGER.warn( - `du ${path} in ${containerName} did not finish within ${timeoutMs / 1000}s` - ) - return null + } catch (error) { + if (!timedOut) { + throw error } } finally { clearTimeout(timer) } + if (timedOut) { + CORE_LOGGER.warn( + `du ${path} in ${containerName} did not finish within ${timeoutMs / 1000}s` + ) + return null + } const output = Buffer.concat(chunks).toString() - // `du -sb` prints "\t", after any "du: cannot access" lines (e.g. for /proc), - // which the TTY merges into the same output. + // The size line follows any "du: cannot access" lines, which the TTY merges in. const match = output.match(/^(\d+)\t/m) return match ? parseInt(match[1], 10) : null } catch (error) { @@ -4879,8 +4873,7 @@ export class C2DEngineDocker extends C2DEngine { ? undefined : await this.sampleModelDownload(job, engine) - // A service that never answers is usually a node that cannot reach its container: say so once, - // where an operator looks, rather than only through users stuck on "warming up". + // Usually a node that cannot reach its container: tell the operator once per container. if ( readiness.state === 'waiting' && Number.isFinite(startedAt) && @@ -4897,24 +4890,7 @@ export class C2DEngineDocker extends C2DEngine { ) } - // Same cross-process guard as the metrics write: a lifecycle transition must win. - if ( - this.serviceOpsInFlight.has(job.serviceId) || - (await this.db.isServiceLocked(job.serviceId, SERVICE_LOCK_STALE_MS)) - ) { - return - } - await this.db.updateServiceJobReadiness( - job.serviceId, - { - owner: job.owner, - clusterHash: job.clusterHash, - status: ServiceStatusNumber.Running, - containerId: job.containerId - }, - readiness, - modelDownload - ) + await this.persistServiceReadiness(job, readiness, modelDownload) } catch (e: any) { CORE_LOGGER.debug( `[readiness] service ${job.serviceId}: probe failed: ${e?.message}` @@ -4922,14 +4898,35 @@ export class C2DEngineDocker extends C2DEngine { } } + // Same cross-process guard as the metrics write: a lifecycle transition must win. + private async persistServiceReadiness( + job: ServiceJob, + readiness: ServiceReadiness | undefined, + modelDownload: ServiceModelDownload | undefined + ): Promise { + if ( + this.serviceOpsInFlight.has(job.serviceId) || + (await this.db.isServiceLocked(job.serviceId, SERVICE_LOCK_STALE_MS)) + ) { + return + } + await this.db.updateServiceJobReadiness( + job.serviceId, + { + owner: job.owner, + clusterHash: job.clusterHash, + status: ServiceStatusNumber.Running, + containerId: job.containerId + }, + readiness, + modelDownload + ) + } + /** - * Download progress for a service the node cannot probe: only what the launch script's download - * manifest says, with no readiness, so Running stays the signal that the service is usable. - * - * Throttled like the probe. A service with no manifest MANIFEST_DISCOVERY_MS after its container - * started is not checked again, and sampling stops once the listed download is complete or after - * MANIFEST_SAMPLING_MS, whichever comes first — with no - * readiness to end it, it would otherwise run for the whole session. + * Manifest download progress for a service the node cannot probe; no readiness, so Running stays + * the signal. With no readiness to end it, sampling is bounded by MANIFEST_DISCOVERY_MS / + * MANIFEST_SAMPLING_MS and stops once the listed download is complete. */ private async sampleUnprobedServiceDownload( job: ServiceJob, @@ -4960,42 +4957,26 @@ export class C2DEngineDocker extends C2DEngine { this.docker.getContainer(job.containerId), (path) => this.manifestDirectoryBytes(job, path) ) - if (!modelDownload) { - return - } - // Same cross-process guard as the readiness write: a lifecycle transition must win. - if ( - this.serviceOpsInFlight.has(job.serviceId) || - (await this.db.isServiceLocked(job.serviceId, SERVICE_LOCK_STALE_MS)) - ) { - return + if (modelDownload) { + await this.persistServiceReadiness(job, undefined, modelDownload) } - await this.db.updateServiceJobReadiness( - job.serviceId, - { - owner: job.owner, - clusterHash: job.clusterHash, - status: ServiceStatusNumber.Running, - containerId: job.containerId - }, - undefined, - modelDownload - ) } - /** A manifest folder's size, re-measured with `du` at most every MANIFEST_DIR_MEASURE_MS. */ private async manifestDirectoryBytes( job: ServiceJob, path: string ): Promise { - let cache = this.manifestDirSizes.get(job.serviceId) - if (!cache || cache.containerId !== job.containerId) { - cache = { containerId: job.containerId, sizes: new Map() } - this.manifestDirSizes.set(job.serviceId, cache) + let sizes = this.manifestDirSizes.get(job.serviceId) + if (!sizes) { + sizes = new Map() + this.manifestDirSizes.set(job.serviceId, sizes) } const now = Date.now() - const last = cache.sizes.get(path) - if (last && now - last.at < MANIFEST_DIR_MEASURE_MS) { + const last = sizes.get(path) + if ( + last?.containerId === job.containerId && + now - last.at < MANIFEST_DIR_MEASURE_MS + ) { return last.bytes } const bytes = await this.getContainerDiskUsage( @@ -5003,7 +4984,7 @@ export class C2DEngineDocker extends C2DEngine { path, MANIFEST_DU_TIMEOUT_MS ) - cache.sizes.set(path, { bytes, at: now }) + sizes.set(path, { containerId: job.containerId, bytes, at: now }) return bytes } @@ -5025,8 +5006,7 @@ export class C2DEngineDocker extends C2DEngine { ): Promise { try { const container = this.docker.getContainer(job.containerId) - // The launch script's own manifest, when it writes one (see downloadManifest). Read before - // the completion check below: a script can append to it after its first downloads finish. + // Before the completion check: a script can append to its manifest after earlier downloads. const fromManifest = await sampleDownloadManifest(container, (path) => this.manifestDirectoryBytes(job, path) ) diff --git a/src/components/c2d/downloadManifest.ts b/src/components/c2d/downloadManifest.ts index 555f51159..b25c06648 100644 --- a/src/components/c2d/downloadManifest.ts +++ b/src/components/c2d/downloadManifest.ts @@ -5,29 +5,22 @@ import type { ServiceModelDownload } from '../../@types/C2D/ServiceOnDemand.js' import { buildModelDownload, fileSize } from './modelDownload.js' /** - * Download progress from a manifest the service's own launch script writes. + * Download progress from a manifest the service's launch script writes before its first download, + * one `\t\t[\t]` line per item (`bytes` 0 = unknown): * - * Before its first download, a template script lists what it will fetch in MANIFEST_PATH, one - * `\t\t[\t]` line per item: + * file written to `.part`, renamed to `` once complete + * dir a folder filled under names of its own, measured with `du` * - * file one file, written to `.part` and renamed to `` once complete - * dir a folder a tool fills under names of its own (huggingface_hub's snapshot cache), - * measured with `du` and complete once it holds `` - * - * `bytes` is what the item will weigh, or 0 when the script could not tell. A script may append - * lines while it runs, so the file is re-read on every sample. It lives in the container's own - * filesystem, so a new container never sees an earlier one's. Like the script that writes it, the - * manifest is untrusted: it only ever decides what this service's own progress report says. + * Re-read on every sample, since the script may append to it. Untrusted: it only affects this + * service's own progress report. */ export const MANIFEST_PATH = '/tmp/ocean/downloads.tsv' const MAX_DOWNLOAD_ENTRIES = 64 -// Each `dir` entry costs a `du` inside the container, so far fewer of them are measured. +// Each `dir` entry costs a `du` inside the container. const MAX_DIR_ENTRIES = 8 -// How long after its container starts a service the node cannot probe is checked for a manifest. -// A script writes it before its first download; setup steps before that can take minutes. +// For a service the node cannot probe: how long after start to look for a manifest, and the +// longest it is sampled, so a download that never completes is not sampled forever. export const MANIFEST_DISCOVERY_MS = 30 * 60 * 1000 -// The longest a service the node cannot probe is sampled after its container starts, so a download -// that never reads as complete (an unknown size, a file that never arrives) is not sampled forever. export const MANIFEST_SAMPLING_MS = 6 * 60 * 60 * 1000 const MAX_READ_BYTES = 64 * 1024 const BYTES = /^\d{1,13}$/ @@ -38,10 +31,8 @@ export interface DownloadEntry { bytes: number | null } -/** Bytes held under a folder in the container, or null when it is not there yet. */ export type DirectoryBytes = (path: string) => Promise -/** One small file from the container, via Docker's archive endpoint. Null when absent or too big. */ async function readContainerFile( container: Dockerode.Container, path: string @@ -66,10 +57,7 @@ async function readContainerFile( } } -/** - * The manifest's entries, capped at MAX_DOWNLOAD_ENTRIES (MAX_DIR_ENTRIES of them `dir`). `truncated` - * says it names more than that, so the entries are not the whole download. - */ +/** `truncated`: the manifest names more than the caps allow, so it must never read as complete. */ export function parseDownloadManifest(text: string): { entries: DownloadEntry[] truncated: boolean @@ -101,26 +89,29 @@ export function parseDownloadManifest(text: string): { return { entries: [...entries.values()], truncated: false } } -/** - * Bytes on disk against the sizes the entries declare. The total and the file count are reported - * only when they are the whole download, so a truncated list never reads as complete. - */ -export async function measureDownloads( +/** Null when the script writes no manifest. */ +export async function sampleDownloadManifest( container: Dockerode.Container, - entries: DownloadEntry[], - truncated: boolean, - directoryBytes?: DirectoryBytes -): Promise { + directoryBytes: DirectoryBytes +): Promise { + const manifest = await readContainerFile(container, MANIFEST_PATH) + if (!manifest) { + return null + } + const { entries, truncated } = parseDownloadManifest(manifest) + if (entries.length === 0) { + return null + } let downloadedBytes = 0 let files = 0 let inFlight = 0 for (const { kind, path, bytes } of entries) { if (kind === 'dir') { - const held = directoryBytes ? await directoryBytes(path) : null + const held = await directoryBytes(path) if (held === null) { continue } - // Capped: a reused cache can hold more than this launch asked for. + // A reused cache can hold more than this launch asked for. downloadedBytes += bytes ? Math.min(held, bytes) : held if (bytes && held >= bytes) { files++ @@ -129,7 +120,7 @@ export async function measureDownloads( } continue } - // Finished first, so a stale .part left beside it cannot hold the file in flight forever. + // Finished first, so a stale .part beside it cannot keep the file in flight. const finished = await fileSize(container, path) if (finished !== null) { downloadedBytes += finished @@ -151,23 +142,3 @@ export async function measureDownloads( ...(truncated ? {} : { filesTotal: entries.length }) } } - -/** Progress from the launch script's manifest. Null when the script writes none. Never throws. */ -export async function sampleDownloadManifest( - container: Dockerode.Container, - directoryBytes: DirectoryBytes -): Promise { - const manifest = await readContainerFile(container, MANIFEST_PATH) - if (!manifest) { - return null - } - const { entries, truncated } = parseDownloadManifest(manifest) - if (entries.length === 0) { - return null - } - try { - return await measureDownloads(container, entries, truncated, directoryBytes) - } catch { - return null - } -} diff --git a/src/components/c2d/modelDownload.ts b/src/components/c2d/modelDownload.ts index d8b33fc56..eba6801a1 100644 --- a/src/components/c2d/modelDownload.ts +++ b/src/components/c2d/modelDownload.ts @@ -470,8 +470,8 @@ async function fetchGgufFileBytes( * Whether a recorded download has finished: the total reached AND nothing still arriving. The total * is an estimate (dtype × parameters), so on its own it can read 100% while the last shard is still * being written; no partial files left is what makes it final. A record with no total never - * completes here, and sampling simply carries on until the service is ready. A record that knows - * its file count (a download manifest) is complete once every listed item is. + * completes here, and sampling simply carries on until the service is ready. A manifest record is + * complete once every listed item is. */ export function isModelDownloadComplete( download: ServiceModelDownload | undefined diff --git a/src/components/c2d/serviceEngines.ts b/src/components/c2d/serviceEngines.ts index 21013eb87..7cd7cfd0b 100644 --- a/src/components/c2d/serviceEngines.ts +++ b/src/components/c2d/serviceEngines.ts @@ -153,8 +153,7 @@ export const SERVICE_ENGINE_PROFILES: ServiceEngineProfile[] = [ { id: 'comfyui', matchesImage: (image) => /(^|\/)yanwk\/comfyui-boot$/.test(image), - // ComfyUI binds its port only after the bundle script has fetched its models and every custom - // node has loaded, so the first 200 means the UI is usable. + // Binds its port only after models are fetched and custom nodes loaded. probe: { path: '/system_stats', port: 8188, expectStatus: [200] } } ] @@ -173,10 +172,8 @@ function stripImageRef(image: string): string { } /** - * The engine profile for a service, or null when the node does not recognize the image — or when - * the service does not publish the engine's port. An app that runs the engine behind itself (vLLM - * on loopback, the app on its own port) is not the engine's to call ready: probing the app's port - * for the engine's path would never succeed. + * Null for an unrecognized image, or when the engine's port is not published (an app running the + * engine behind its own port has nothing to probe). */ export function resolveServiceEngine(job: ServiceJob): ServiceEngineProfile | null { const image = stripImageRef((job.image || job.containerImage || '').trim()) diff --git a/src/components/c2d/serviceReadiness.ts b/src/components/c2d/serviceReadiness.ts index 8c082cc3a..e775c3a0e 100644 --- a/src/components/c2d/serviceReadiness.ts +++ b/src/components/c2d/serviceReadiness.ts @@ -111,8 +111,7 @@ export const PROBE_PERIOD_SECONDS = 5 // Once ready the check keeps running — to catch an engine that dies without its container exiting — // but far more slowly: it is a liveness check at that point, not a wait. export const READY_PROBE_PERIOD_SECONDS = 30 -// A service still not ready this long after its container started is logged as a warning, once per -// container: usually a node that cannot reach it, sometimes just a very large download. +// Not ready this long after start: warn once per container. export const READINESS_WARN_AFTER_SECONDS = 15 * 60 const PROBE_TIMEOUT_MS = 2000 diff --git a/src/components/database/sqliteCompute.ts b/src/components/database/sqliteCompute.ts index a801cf2f4..314623f43 100644 --- a/src/components/database/sqliteCompute.ts +++ b/src/components/database/sqliteCompute.ts @@ -409,7 +409,7 @@ export class SQLiteCompute implements ComputeDatabaseProvider { // Same guarantees as updateServiceJobMetrics, for the readiness probe result: it is sampled from // the same lease-free background loop, so it must never overwrite a lifecycle transition either. - // `readiness` is absent for a service the node cannot probe, which reports download progress only. + // `readiness` is undefined for a service the node cannot probe. // eslint-disable-next-line require-await async updateServiceJobReadiness( serviceId: string,