dsh-multimodal-node 0.0.0-stage → 0.3.1

package/lib/pool.js

@@ -0,0 +1,591 @@
+/**
+ * The task pool: everything a client POSTs to the gateway lands here, and the
+ * scheduler decides which bound instance runs it.
+ *
+ * The contract with callers is ComfyUI's own: POST /prompt returns a prompt id
+ * immediately, and the caller polls /history/<id> until it appears. The pool
+ * therefore owns identifiers, not the backends — one gateway prompt id maps to
+ * one upstream prompt id on whichever instance got the work.
+ *
+ * Two orderings, switched at runtime:
+ *
+ *   fifo     — strictly by arrival (and explicit priority), what a plain queue
+ *              does. Predictable, and what a user expects from "按先后顺序来".
+ *   dynamic  — "least total time": every task gets a cost estimate from its own
+ *              graph, every instance keeps an EWMA of seconds-per-cost-unit, and
+ *              the dispatcher repeatedly assigns the most expensive ready task
+ *              to the instance that would finish it soonest (LPT over a
+ *              machine-aware time model). Explicit priority still jumps the
+ *              line, so a user can force a task in front.
+ */
+import { randomUUID } from 'node:crypto';
+import { Backend } from './comfy.js';
+
+const MAX_ATTEMPTS = 2;
+const TICK_MS = 400;
+
+/** Rough "cost units" for one prompt graph: sampler work dominates, the rest is a floor. */
+export function estimateCost(prompt) {
+    const nodes = prompt !== null && typeof prompt === 'object' ? prompt : {};
+    const byId = new Map(Object.entries(nodes));
+    const resolved = new Map();
+
+    /** Follow one input link back to the node that produces it (bounded depth). */
+    const follow = (value, depth = 0) => {
+        if (depth > 6 || !Array.isArray(value) || value.length < 1) return undefined;
+        const node = byId.get(String(value[0]));
+        if (node === undefined) return undefined;
+        return node;
+    };
+
+    const latentSize = (node) => {
+        if (node === undefined || typeof node !== 'object') return { mp: 1, batch: 1 };
+        const inputs = node.inputs ?? {};
+        if (inputs.width !== undefined && inputs.height !== undefined) {
+            const batch = Number(inputs.batch_size ?? 1) || 1;
+            return { mp: (Number(inputs.width) * Number(inputs.height)) / 1e6, batch };
+        }
+        if (Array.isArray(inputs.latent_image) || Array.isArray(inputs.samples)) {
+            return latentSize(follow(inputs.latent_image ?? inputs.samples, 1));
+        }
+        if (Array.isArray(inputs.pixels)) {
+            const source = follow(inputs.pixels);
+            return source === undefined ? { mp: 1, batch: 1 } : latentSize(source);
+        }
+        return { mp: 1, batch: 1 };
+    };
+
+    let cost = 0;
+    let samplers = 0;
+    for (const [id, node] of byId) {
+        if (node === null || typeof node !== 'object') continue;
+        const type = String(node.class_type ?? '');
+        const inputs = node.inputs ?? {};
+        if (/sampler/i.test(type) && inputs.steps !== undefined) {
+            samplers += 1;
+            const steps = Number(inputs.steps) || 20;
+            const { mp, batch } = latentSize({ inputs, class_type: type });
+            const cfg = Number(inputs.cfg ?? 7) || 7;
+            // CFG batches the UNet twice above 1.0 — the single biggest term.
+            const cfgFactor = cfg > 1.01 ? 2 : 1;
+            cost += steps * mp * batch * cfgFactor;
+            resolved.set(id, cost);
+        }
+    }
+    // VAE decode/encode and upscales cost real time but not sampler steps.
+    for (const [id, node] of byId) {
+        if (node === null || typeof node !== 'object') continue;
+        const type = String(node.class_type ?? '');
+        if (/VAEDecode|VAEEncode|Upscale|ImageScale/i.test(type)) {
+            const { mp, batch } = latentSize({ inputs: node.inputs ?? {}, class_type: type });
+            cost += 0.5 * mp * batch;
+            resolved.set(id, cost);
+        }
+    }
+    const nodeCount = byId.size;
+    // Model load: the first run after an unload pays a VRAM fill; approximate.
+    cost += nodeCount * 0.05;
+    if (samplers === 0) cost += 1;
+    return Math.max(0.1, cost);
+}
+
+export class TaskPool {
+    /**
+     * @param {{instances: import('./instances.js').InstanceManager, settings: () => any,
+     *          onChange: () => void, log?: (level: string, message: string) => void}} options
+     */
+    constructor(options) {
+        this.options = options;
+        this.tasks = new Map(); // gateway id -> task
+        this.order = []; // ids of live tasks, kept sorted on demand
+        this.timer = null;
+        this.ticking = false;
+        this.listeners = new Set();
+        this.counters = { submitted: 0, completed: 0, failed: 0, retried: 0 };
+        this.lastTickAt = 0;
+        /** instanceId -> { by, token } while a person holds that canvas. */
+        this.bridgeLocks = new Map();
+    }
+
+    /**
+     * An instance being driven by a person (the remote console's canvas) leaves
+     * the rotation: two independent producers queueing on one GPU is how VRAM
+     * runs out mid-sampling. Passing `null` hands it back.
+     */
+    setBridgeLock(instanceId, lock) {
+        if (lock === null || lock === undefined) this.bridgeLocks.delete(instanceId);
+        else this.bridgeLocks.set(instanceId, lock);
+        this.notify();
+        void this.pump();
+    }
+
+    bridgeLockOf(instanceId) {
+        return this.bridgeLocks.get(instanceId) ?? null;
+    }
+
+    /**
+     * Move one queued task in the visible order. The pool is order-driven rather
+     * than slot-driven, so "up" means "priority above the task before me" — the
+     * only thing a person can actually mean by dragging a row.
+     */
+    reorder(id, delta) {
+        const task = this.tasks.get(id);
+        if (task === undefined) return { ok: false, error: `unknown prompt ${id}` };
+        if (task.state !== 'queued') return { ok: false, error: `只有排队中的任务能调整顺序（当前 ${task.state}）` };
+        if (!Number.isFinite(delta) || delta === 0) return { ok: false, error: 'delta 不能为 0' };
+        return this.prioritize(id, delta > 0 ? -1 : 1);
+    }
+
+    get instances() {
+        return this.options.instances;
+    }
+
+    toggleDynamic(enabled) {
+        this.options.settings().dynamic = enabled === true;
+        this.options.onChange();
+        this.notify();
+        void this.pump();
+    }
+
+    notify() {
+        for (const fn of this.listeners) {
+            try {
+                fn();
+            } catch {
+                // a broken listener must not stop scheduling
+            }
+        }
+    }
+
+    subscribe(fn) {
+        this.listeners.add(fn);
+        return () => this.listeners.delete(fn);
+    }
+
+    backend(runtime) {
+        return new Backend({
+            id: runtime.id,
+            baseUrl: runtime.instance.baseUrl,
+            apiKey: runtime.instance.apiKey,
+        });
+    }
+
+    /** Queue one graph. Returns the gateway prompt id at once. */
+    enqueue(body) {
+        const prompt = body?.prompt;
+        if (prompt === null || typeof prompt !== 'object' || Array.isArray(prompt) || Object.keys(prompt).length === 0) {
+            throw new Error('prompt 必须是节点对象（ComfyUI /prompt 的 body.prompt）');
+        }
+        const requestedPriority = Number(
+            body.priority ?? body.extra_data?.gateway_priority ?? body.extra_data?.priority ?? 0,
+        );
+        const id = typeof body.prompt_id === 'string' && body.prompt_id !== '' ? body.prompt_id : randomUUID();
+        const task = {
+            id,
+            seq: this.counters.submitted + 1,
+            prompt,
+            clientId: typeof body.client_id === 'string' ? body.client_id : '',
+            extraData: body.extra_data ?? {},
+            priority: Number.isFinite(requestedPriority) ? Math.max(-100, Math.min(100, Math.floor(requestedPriority))) : 0,
+            cost: estimateCost(prompt),
+            enqueuedAt: Date.now(),
+            startedAt: null,
+            finishedAt: null,
+            state: 'queued',
+            backendId: null,
+            upstreamId: null,
+            attempts: 0,
+            error: null,
+            outputs: null,
+            historyEntry: null,
+            progress: null,
+            pinned: typeof body.endpoint === 'string' && body.endpoint !== '' && body.endpoint !== 'auto' ? body.endpoint : null,
+        };
+        this.tasks.set(id, task);
+        this.order.push(id);
+        this.counters.submitted += 1;
+        this.options.log?.('info', `enqueue ${id.slice(0, 8)} cost=${task.cost.toFixed(1)}`);
+        this.notify();
+        void this.pump();
+        return { prompt_id: id, number: task.seq, node_errors: {} };
+    }
+
+    /** Ready tasks in dispatch order for the current policy. */
+    ordered() {
+        const live = this.order
+            .map((id) => this.tasks.get(id))
+            .filter((task) => task !== undefined && task.state === 'queued');
+        const dynamic = this.options.settings().dynamic === true;
+        live.sort((left, right) => {
+            if (left.priority !== right.priority) return right.priority - left.priority;
+            if (dynamic) {
+                // LPT: the most expensive ready task goes first, so a fast GPU is
+                // never left holding a cheap job while a big one waits.
+                if (Math.abs(left.cost - right.cost) > 1e-6) return right.cost - left.cost;
+            }
+            return left.enqueuedAt - right.enqueuedAt;
+        });
+        return live;
+    }
+
+    /** Free capacity for one instance: maxConcurrent minus what it is running. */
+    freeSlots(runtime) {
+        const running = [...this.tasks.values()].filter(
+            (task) => task.state === 'running' && task.backendId === runtime.id,
+        ).length;
+        return runtime.instance.maxConcurrent - running;
+    }
+
+    pickBackend(task) {
+        const candidates = this.instances
+            .enabled()
+            .filter((runtime) => runtime.health.ok !== false && !runtime.starting && !runtime.stopping)
+            .filter((runtime) => !this.bridgeLocks.has(runtime.id))
+            .filter((runtime) => this.freeSlots(runtime) > 0);
+        if (candidates.length === 0) return undefined;
+        if (task.pinned !== null) {
+            const pinned = candidates.find((runtime) => runtime.id === task.pinned);
+            if (pinned !== undefined) return pinned;
+        }
+        // Time model: seconds = speed × cost. Prefer the instance that finishes
+        // this task soonest, weighted by the operator's dispatch weight.
+        let best;
+        let bestScore = Infinity;
+        for (const runtime of candidates) {
+            const speed = runtime.speed > 0 ? runtime.speed : 1;
+            const queueAhead = [...this.tasks.values()].filter(
+                (other) => other.state === 'running' && other.backendId === runtime.id,
+            ).reduce((sum, other) => sum + other.cost * speed, 0);
+            const score = (queueAhead + task.cost * speed) / Math.max(1, runtime.instance.weight);
+            if (score < bestScore) {
+                bestScore = score;
+                best = runtime;
+            }
+        }
+        return best;
+    }
+
+    /** One scheduling pass: fill every free slot it can, then arm the timer. */
+    async pump() {
+        if (this.ticking) return;
+        this.ticking = true;
+        this.lastTickAt = Date.now();
+        try {
+            for (;;) {
+                const queue = this.ordered();
+                if (queue.length === 0) break;
+                const task = queue[0];
+                const runtime = this.pickBackend(task);
+                if (runtime === undefined) break;
+                queue.shift();
+                await this.dispatch(task, runtime);
+            }
+        } finally {
+            this.ticking = false;
+            this.notify();
+        }
+    }
+
+    async dispatch(task, runtime) {
+        task.state = 'running';
+        task.startedAt = Date.now();
+        task.backendId = runtime.id;
+        task.attempts += 1;
+        task.error = null;
+        const body = {
+            prompt: task.prompt,
+            client_id: task.clientId !== '' ? task.clientId : `gateway-${runtime.id}`,
+            extra_data: task.extraData ?? {},
+        };
+        try {
+            const reply = await this.backend(runtime).prompt(body);
+            task.upstreamId = typeof reply?.prompt_id === 'string' ? reply.prompt_id : null;
+            this.options.log?.('info', `${task.id.slice(0, 8)} -> ${runtime.id} (${task.upstreamId?.slice(0, 8) ?? '?'})`);
+        } catch (error) {
+            const message = error instanceof Error ? error.message : String(error);
+            task.error = message;
+            runtime.lastError = message;
+            runtime.health = { ...runtime.health, ok: false, checkedAt: Date.now(), error: message };
+            this.counters.retried += 1;
+            if (task.attempts < MAX_ATTEMPTS) {
+                task.state = 'queued';
+                task.backendId = null;
+                task.upstreamId = null;
+                task.startedAt = null;
+                this.options.log?.('warn', `${task.id.slice(0, 8)} retry after: ${message}`);
+            } else {
+                task.state = 'failed';
+                task.finishedAt = Date.now();
+                this.counters.failed += 1;
+            }
+        }
+        this.notify();
+    }
+
+    /** Poll the backends: finished runs, dropped runs, progress. */
+    async collect() {
+        const running = [...this.tasks.values()].filter((task) => task.state === 'running' && task.upstreamId !== null);
+        const byBackend = new Map();
+        for (const task of running) {
+            const list = byBackend.get(task.backendId) ?? [];
+            list.push(task);
+            byBackend.set(task.backendId, list);
+        }
+        await Promise.all(
+            [...byBackend].map(async ([backendId, tasks]) => {
+                const runtime = this.instances.get(backendId);
+                if (runtime === undefined) return;
+                const client = this.backend(runtime);
+                let queue;
+                try {
+                    queue = await client.queue();
+                } catch {
+                    return; // a dead instance is handled by its own health probe
+                }
+                const runningIds = new Set((queue.queue_running ?? []).map((entry) => entry[1]));
+                const pendingIds = new Set((queue.queue_pending ?? []).map((entry) => entry[1]));
+                for (const task of tasks) {
+                    const position = runningIds.has(task.upstreamId) ? 'running' : pendingIds.has(task.upstreamId) ? 'pending' : 'gone';
+                    task.progress = { position };
+                    if (position !== 'gone') continue;
+                    let entry;
+                    try {
+                        const history = await client.historyOf(task.upstreamId);
+                        entry = history?.[task.upstreamId];
+                    } catch {
+                        entry = undefined;
+                    }
+                    if (entry === undefined) {
+                        // Not queued and not in history: the instance lost it (restart,
+                        // cleartext cancel, crash). Retry elsewhere rather than hang.
+                        if (task.attempts < MAX_ATTEMPTS) {
+                            task.state = 'queued';
+                            task.backendId = null;
+                            task.upstreamId = null;
+                            task.startedAt = null;
+                            this.counters.retried += 1;
+                            this.options.log?.('warn', `${task.id.slice(0, 8)} lost on ${backendId}, requeued`);
+                        } else {
+                            task.state = 'failed';
+                            task.finishedAt = Date.now();
+                            task.error = '任务在执行端消失（实例重启或被清队）';
+                            this.counters.failed += 1;
+                        }
+                        continue;
+                    }
+                    this.finish(task, runtime, entry);
+                }
+            }),
+        );
+    }
+
+    finish(task, runtime, entry) {
+        const status = entry?.status ?? {};
+        const completed = status.completed === true || status.status_str === 'success';
+        task.historyEntry = entry;
+        task.outputs = entry?.outputs ?? {};
+        task.finishedAt = Date.now();
+        task.state = completed ? 'done' : 'failed';
+        task.error = completed ? null : (status.status_str === 'error' ? '执行失败' : status.status_str ?? '未知状态');
+        if (!completed) this.counters.failed += 1;
+        else this.counters.completed += 1;
+        // Learn this instance's speed: seconds per cost unit, EWMA over the last
+        // runs. One slow run (a cold model load) must not skew the scheduler.
+        const seconds = task.startedAt === null ? null : (task.finishedAt - task.startedAt) / 1000;
+        if (seconds !== null && seconds > 0.2 && task.cost > 0) {
+            const sample = seconds / task.cost;
+            runtime.speed = runtime.speedSamples === 0 ? sample : runtime.speed * 0.7 + sample * 0.3;
+            runtime.speedSamples += 1;
+        }
+        this.options.log?.(
+            completed ? 'info' : 'error',
+            `${task.id.slice(0, 8)} ${completed ? 'done' : 'failed'} on ${runtime.id} in ${seconds === null ? '?' : seconds.toFixed(1)}s`,
+        );
+        this.trim();
+        this.notify();
+    }
+
+    /** Keep finished tasks for the history view, but bounded. */
+    trim() {
+        const keepFinished = 200;
+        const finished = this.order
+            .map((id) => this.tasks.get(id))
+            .filter((task) => task !== undefined && task.state !== 'queued' && task.state !== 'running');
+        const excess = finished.length - keepFinished;
+        if (excess <= 0) return;
+        for (const task of finished.slice(0, excess)) {
+            this.tasks.delete(task.id);
+        }
+        this.order = this.order.filter((id) => this.tasks.has(id));
+    }
+
+    cancel(id) {
+        const task = this.tasks.get(id);
+        if (task === undefined) return { ok: false, error: `unknown prompt ${id}` };
+        if (task.state === 'queued') {
+            task.state = 'cancelled';
+            task.finishedAt = Date.now();
+            this.notify();
+            return { ok: true, cancelled: 'queued' };
+        }
+        if (task.state === 'running') {
+            const runtime = task.backendId === null ? undefined : this.instances.get(task.backendId);
+            task.state = 'cancelled';
+            task.finishedAt = Date.now();
+            if (runtime !== undefined && task.upstreamId !== null) {
+                // Pending: dequeue; running: interrupt only this prompt.
+                const client = this.backend(runtime);
+                void client
+                    .deleteQueueItems([task.upstreamId])
+                    .catch(() => {})
+                    .then(() => client.interrupt(task.upstreamId).catch(() => {}));
+            }
+            this.notify();
+            return { ok: true, cancelled: 'running' };
+        }
+        return { ok: false, error: `prompt ${id} 已结束（${task.state}）` };
+    }
+
+    /** Explicit priority bump: the "插队" button in the UI and the tool. */
+    prioritize(id, delta) {
+        const task = this.tasks.get(id);
+        if (task === undefined) return { ok: false, error: `unknown prompt ${id}` };
+        if (task.state !== 'queued') return { ok: false, error: `只有排队中的任务能调整优先级（当前 ${task.state}）` };
+        task.priority = Math.max(-100, Math.min(100, task.priority + Math.floor(Number(delta) || 1)));
+        this.notify();
+        void this.pump();
+        return { ok: true, priority: task.priority };
+    }
+
+    /** Sum of `cost × speed` still ahead: a lower bound on remaining wall time. */
+    estimateRemaining() {
+        const now = Date.now();
+        const perInstance = new Map();
+        for (const runtime of this.instances.enabled()) perInstance.set(runtime.id, { busyUntil: 0, speed: runtime.speed || 1, id: runtime.id });
+        const work = [...this.tasks.values()].filter((task) => task.state === 'queued' || task.state === 'running');
+        // Running work already occupies its instance.
+        for (const task of work.filter((task) => task.state === 'running')) {
+            const slot = perInstance.get(task.backendId);
+            if (slot !== undefined) slot.busyUntil += task.cost * slot.speed * 1000;
+        }
+        for (const task of this.ordered()) {
+            let best;
+            for (const slot of perInstance.values()) {
+                if (best === undefined || slot.busyUntil < best.busyUntil) best = slot;
+            }
+            if (best === undefined) break;
+            best.busyUntil += task.cost * best.speed * 1000;
+        }
+        const finishes = [...perInstance.values()].map((slot) => slot.busyUntil);
+        const maxFinish = finishes.length === 0 ? 0 : Math.max(...finishes);
+        return { makespanMs: Math.round(maxFinish), at: now };
+    }
+
+    /** Merge our view with what the instances report about foreign work. */
+    async queueView() {
+        const running = [];
+        const pending = this.ordered().map((task, index) => [
+            task.seq,
+            task.id,
+            task.prompt,
+            task.extraData,
+            Object.keys(task.prompt),
+            index,
+        ]);
+        for (const task of this.tasks.values()) {
+            if (task.state !== 'running') continue;
+            running.push([task.seq, task.upstreamId ?? task.id, task.prompt, task.extraData, Object.keys(task.prompt), task.progress?.position === 'running' ? 0 : 0]);
+        }
+        // Foreign work already sitting in an instance's own queue is real and must
+        // stay visible, otherwise the frontend believes the GPU is free.
+        await Promise.all(
+            this.instances.enabled().map(async (runtime) => {
+                try {
+                    const queue = await this.backend(runtime).queue();
+                    for (const entry of queue.queue_running ?? []) {
+                        if (entry[1] === undefined || [...this.tasks.values()].some((task) => task.upstreamId === entry[1])) continue;
+                        running.push(entry);
+                    }
+                    for (const entry of queue.queue_pending ?? []) {
+                        if (entry[1] === undefined || [...this.tasks.values()].some((task) => task.upstreamId === entry[1])) continue;
+                        pending.push(entry);
+                    }
+                } catch {
+                    // an unreachable instance simply contributes nothing
+                }
+            }),
+        );
+        return { queue_running: running, queue_pending: pending };
+    }
+
+    snapshot() {
+        const tasks = this.order
+            .map((id) => this.tasks.get(id))
+            .filter((task) => task !== undefined)
+            .slice()
+            .sort((left, right) => {
+                const rank = (task) => (task.state === 'running' ? 0 : task.state === 'queued' ? 1 : 2);
+                if (rank(left) !== rank(right)) return rank(left) - rank(right);
+                if (left.state === 'queued' && right.state === 'queued') return (right.priority - left.priority) || (left.enqueuedAt - right.enqueuedAt);
+                return (right.finishedAt ?? right.enqueuedAt) - (left.finishedAt ?? left.enqueuedAt);
+            });
+        return {
+            dynamic: this.options.settings().dynamic === true,
+            counters: this.counters,
+            estimate: this.estimateRemaining(),
+            tasks: tasks.map((task) => ({
+                id: task.id,
+                seq: task.seq,
+                state: task.state,
+                backendId: task.backendId,
+                priority: task.priority,
+                cost: Number(task.cost.toFixed(2)),
+                attempts: task.attempts,
+                enqueuedAt: task.enqueuedAt,
+                startedAt: task.startedAt,
+                finishedAt: task.finishedAt,
+                error: task.error,
+                progress: task.progress,
+                pinned: task.pinned,
+                workflow: task.extraData?.workflowName ?? task.extraData?.workflow ?? null,
+                outputs: countOutputs(task.outputs),
+            })),
+        };
+    }
+
+    start() {
+        if (this.timer !== null) return;
+        this.timer = setInterval(() => {
+            void (async () => {
+                try {
+                    await this.instances.probeAll(false);
+                    await this.collect();
+                    await this.pump();
+                } catch {
+                    // a scheduling pass never throws into the timer
+                }
+            })();
+        }, TICK_MS);
+        if (typeof this.timer.unref === 'function') this.timer.unref();
+    }
+
+    dispose() {
+        if (this.timer !== null) {
+            clearInterval(this.timer);
+            this.timer = null;
+        }
+        this.listeners.clear();
+    }
+}
+
+function countOutputs(outputs) {
+    if (outputs === null || typeof outputs !== 'object') return 0;
+    let total = 0;
+    for (const value of Object.values(outputs)) {
+        if (value === null || typeof value !== 'object') continue;
+        for (const list of Object.values(value)) {
+            if (Array.isArray(list)) total += list.length;
+        }
+    }
+    return total;
+}