From 3b9345f331622f6a4634a1fdc6281fecdc2e028e Mon Sep 17 00:00:00 2001 From: "[._.]/ Adam Eivy" Date: Sat, 15 Aug 2026 08:51:40 +0000 Subject: [PATCH 1/3] feat: serialize heavy local accelerator jobs (#4174) --- .changelog/next/added-issue-4174.md | 1 + server/lib/README.md | 2 + server/lib/heavyJobClaim.js | 145 ++++++++++++++++++ server/lib/heavyJobClaim.test.js | 69 +++++++++ server/lib/index.js | 2 + server/lib/localMemory.js | 62 ++++++++ server/services/imageGen/local.js | 15 ++ server/services/imageTo3d/models.js | 10 ++ server/services/imageTo3d/models.test.js | 12 ++ server/services/loraTraining/index.js | 20 ++- server/services/loraTraining/memoryPrep.js | 145 +----------------- .../services/loraTraining/memoryPrep.test.js | 4 +- server/services/videoGen/local.js | 42 ++++- server/services/videoGen/local.test.js | 8 + 14 files changed, 388 insertions(+), 149 deletions(-) create mode 100644 .changelog/next/added-issue-4174.md create mode 100644 server/lib/heavyJobClaim.js create mode 100644 server/lib/heavyJobClaim.test.js create mode 100644 server/lib/localMemory.js diff --git a/.changelog/next/added-issue-4174.md b/.changelog/next/added-issue-4174.md new file mode 100644 index 0000000000..61246f8729 --- /dev/null +++ b/.changelog/next/added-issue-4174.md @@ -0,0 +1 @@ +- Local image, video, 3D, and LoRA jobs now coordinate access to the machine accelerator to prevent overlapping GPU workloads. diff --git a/server/lib/README.md b/server/lib/README.md index 3bf542c1e3..042b58fbbd 100644 --- a/server/lib/README.md +++ b/server/lib/README.md @@ -27,6 +27,7 @@ The barrel `server/lib/index.js` is a machine-checkable enumeration of every pub | `sharedSchemas.js` | Cross-domain Zod fragments that `validation.js` and the per-domain `*Validation.js` files both need, kept in a leaf module so the domain files never import back through `validation.js` (its hoisted `export * from` lines would TDZ). `grokVideoDurationSchema` (grok clip-length union), `cloudModelIdString(message)` (cloud-CLI model-id charset/bounds), `recordRenderPinFields` (the `imageMode`/`imageModelId` per-record render pin pair), and `isSafeSubdirFilter(v)` (relative path with no wildcard, `..` segment, or leading `/`), and `csvIdsParam({ max, maxIdLength, truncate })` (the `?ids=a,b,c` batch-by-id query param: trims, drops empties, reads all-blank as absent, and either 400s or silently slices an over-cap batch). | | `agentOutputMarkers.js` | Status lines PortOS itself appends to an agent's output buffer: `SENTINEL_COMPLETION_MARKER` (the line `ingestDoneSentinel` writes just before the agent's `.agent-done` summary) plus `isAgentLifecycleLine(line)`/`stripLifecycleLines(lines)`, which match PortOS's telemetry on its actual message shapes rather than on a leading emoji (an agent's own summary may well start a line with `✅`). Pure — shared by the emitter (`services/agentTuiSpawning.js`) and by readers that must show only the agent's own words (notably the generated PR description). | | `agentSentinel.js` | The `.agent-done` completion sentinel: `DONE_SENTINEL_NAME`, `doneSentinelName(agentId)` → the per-instance filename `.agent-done-` (worktree-less agents share one workspace, so a shared name lets concurrent runs clobber — and finalize on — each other's signal), `doneSentinelPath(workspacePath, agentId)` → the single path every producer and consumer resolves, + pure `parseSentinelPayload(contents)` → `{ summary, payload }`. Back-compat — a plain-markdown sentinel yields `payload: null`; a JSON object yields its structured `payload` for a programmatic-I/O task type's `processTaskOutput` hook. `salvageSentinelPayload(contents)` (async) is the lenient second tier — runs `jsonExtract` over a fenced/prose-trailed/control-char-corrupted envelope so a less-capable model's near-valid output still surfaces its `payload` instead of being dropped. `extractSentinelPayloadFromTranscript(transcript, isPayload)` (async) is the third tier for programmatic-I/O types ONLY — scans the ANSI-stripped PTY transcript, newest balanced JSON block first, for a payload the model PRINTED instead of writing, and adopts it only if the owning hook's shape predicate accepts it (#3640). | +| `heavyJobClaim.js` | Cross-process, machine-wide accelerator claim: `claimHeavyLocalJob({ kind, id, timeoutMs? })` either acquires an ephemeral `data/` lock or reports its active holder; dead-PID claims are reclaimed, and a spawned child can take over the recorded PID across a server restart. | | `agentValidation.js` | Social-bot agent schemas (personality, Moltbook/Moltworld accounts, automation schedules, agent tools + Moltworld payloads) and CoS Feature Agent definitions. | | `quotaBurnConfig.js` | Quota-burn plan shape: the provider families, the burn-job type alphabet + catalog the config page renders its form from, `QUOTA_BURN_BOUNDS` (the one bounds table the normalizer clamps to, the Zod schemas reject against, and the catalog descriptors publish as min/max), and total normalization (`normalizeQuotaBurnConfig`). Owns the dispatch-cap sentinel too (`QUOTA_BURN_UNLIMITED_DISPATCHES` / `isUnlimitedDispatchCap`) — -1 means the window is not counted, and is the default. Also owns the queued burn task's description shape (`burnTaskDescription` / `quotaBurnFamilyOfDescription`), shared with migration 225, and the `run once` vocabulary (`quotaBurnJobKey` / `jobIsSpent`, plus the two family predicates `familyIsConfigured` and `familyHasRunnableJobs`). Pure — no storage, no provider I/O. | | `quotaBurnPresets.js` | `QUOTA_BURN_PROMPT_PRESETS` — ready-made single-focus audit prompts for `agent-prompt` burn jobs (UX, a11y, mobile, failure paths, perf, test gaps, dead code, data safety, docs, security), each filing GitHub issues and changing no code. Templates: picking one COPIES its prompt into the job, so editing them never rewrites a configured job. `findQuotaBurnPreset(id)`. | @@ -166,6 +167,7 @@ The barrel `server/lib/index.js` is a machine-checkable enumeration of every pub | `imageClean.js` | `cleanImageBuffer(buf, { metadata, denoise })` (composable opt-in pipeline: lossless metadata/C2PA strip + optional median/sharpen denoise) · `stripPngMetadataChunks` / `stripPngC2PAChunk` (lossless PNG-chunk removers) · `compositeIgnoreZone(base, original, mask, { feather })` (preserve-region compositing: restore original pixels into a feathered mask over a diffused result) · `autoCleanGeneratedImage` (in-place clean for post-generation hook). HTTP route in `routes/imageClean.js` wraps `cleanImageBuffer` and appends a CPU light diffusion pass (`applyLightRegen` from `services/imageGen/regen.js`) for the `diffusion=light` SynthID-disruption step. | | `imageWatermark.js` | `removeCornerWatermark` (erases the visible Gemini/Nano-Banana bottom-right ✦ via dependency-free harmonic/Laplace inpaint) + pure helpers `resolveWatermarkRegion` / `inpaintRegion`. Distinct from SynthID regen — this targets the *visible* corner logo. | | `localImageFilename.js` | `localImageFilename(urlOrPath)` resolves a stored image reference to the bare gallery-image filename under `data/images/` (or null for empty/external-URL/non-image-path) — the unit the peer-sync asset pipeline hashes + transfers. Single source of truth for the authors/artists/albums/Creative-Director filename resolvers (`headshotImageFilename`/`portraitImageFilename`/`coverImageFilename`/`startingImageFilename` are thin wrappers). Also exports `assetBasename(pathOrName)`, the shared strip-querystring→basename primitive (reused by moodBoard's `imageUrlToAppAsset`). | +| `localMemory.js` | Shared best-effort local-model unload + available-memory report (`prepareLocalMemory`) for GPU-heavy local work; only loopback Ollama / LM Studio backends are evicted. | | `threejsModel.js` | Validated declarative procedural-model schema plus deterministic standalone Three.js factory export for the Three.js Models workspace. Also the cross-section gate: `evaluateThreejsFlatness(spec)` counts distinct vertex planes per axis — relative to the mesh's own size, with a rotation-invariant zero-volume check behind it — and treats an `extrude` with no bevel thickness as the two-plane slab it is to report when the majority of identity-priority features are built only from flat parts — a model that reads right head-on and like cardboard when orbited. `buildThreejsFlatnessFeedback(flatness)` turns that into default refinement feedback; `listSpecNames(names)` is the shared capped name-list formatter both gates use in finding messages. Also the material-plausibility gate: `evaluateThreejsMaterialPlausibility(spec)` keys a bounded per-family prior table (metal, wood, plastic, glass, fabric, ceramic, rubber, stone, leather, paper) off tokens in each material's id and reports channels whose values the named substance does not support — metalness 0.9 oak, transmission 1.0 steel — skipping any id that names no family or two, and any channel the material's `type` never forwards to Three.js. Advisory only: it never clamps, because a stylized model may legitimately break every prior. `buildThreejsMaterialFeedback(plausibility)` turns that into default refinement feedback. | | `threejsModelCoverage.js` | `evaluateThreejsPartCoverage(spec, { family })` — structural assembly gate over an already-validated Three.js spec: flags promised features fused onto the same part set, geometry claimed by no detail, and details nothing was built for (folded minor relief stays a `note`). With a subject family it also warns on a required component the spec never mentions at all — the one check that can fault a spec for what it failed to *promise*. `buildThreejsCoverageFeedback(coverage)` turns the error findings (plus any family gap) into default refinement feedback. | | `threejsModelFamilies.js` | Curated subject-family checklists for Three.js generation. `THREEJS_MODEL_FAMILY_OPTIONS` / `THREEJS_MODEL_FAMILY_IDS` drive the picker and route validation; `buildThreejsFamilyChecklist(id)` returns the prompt block spliced into generation (`''` for the default `general`, so an un-narrowed subject keeps the unchanged general-purpose prompt); `findMissingFamilyComponents(spec, id)` reports which required components the spec never mentions. | diff --git a/server/lib/heavyJobClaim.js b/server/lib/heavyJobClaim.js new file mode 100644 index 0000000000..7a2a2ca6d1 --- /dev/null +++ b/server/lib/heavyJobClaim.js @@ -0,0 +1,145 @@ +/** + * Machine-wide claim for a local accelerator job. + * + * The claim is deliberately a file below the install's data root rather than an + * in-process flag: PortOS worktrees and restarted server processes share the + * same machine, and must therefore see the same owner. A dead process is + * reclaimed on the next acquisition, so a crash cannot wedge local rendering. + */ + +import { randomUUID } from 'crypto'; +import { mkdir, readFile, unlink, writeFile } from 'fs/promises'; +import { existsSync, readFileSync, unlinkSync } from 'fs'; +import { dirname, join } from 'path'; +import { PATHS } from './fileUtils.js'; + +export const HEAVY_LOCAL_JOB_CLAIM_PATH = join(PATHS.data, 'heavy-local-job.claim.json'); +export const HEAVY_LOCAL_JOB_STALE_MS = 24 * 60 * 60 * 1000; + +const heldClaims = new Map(); +let exitCleanupInstalled = false; + +const sleep = (ms) => new Promise((resolve) => setTimeout(resolve, ms)); + +export function isPidAlive(pid) { + if (!Number.isInteger(pid) || pid <= 0) return false; + try { + process.kill(pid, 0); + return true; + } catch (err) { + return err?.code === 'EPERM'; + } +} + +const parseHolder = (raw) => { + try { + const holder = JSON.parse(raw); + if (!holder || typeof holder.kind !== 'string' || typeof holder.id !== 'string') return null; + return holder; + } catch { + return null; + } +}; + +const readHolder = async (claimPath) => { + const raw = await readFile(claimPath, 'utf8').catch((err) => (err?.code === 'ENOENT' ? null : Promise.reject(err))); + return raw === null ? null : parseHolder(raw); +}; + +const cleanupHeldClaims = () => { + for (const [claimPath, token] of heldClaims) { + try { + if (existsSync(claimPath) && parseHolder(readFileSync(claimPath, 'utf8'))?.token === token) unlinkSync(claimPath); + } catch { + // Process-exit cleanup is best-effort. A stale PID is reclaimed later. + } + } +}; + +const installExitCleanup = () => { + if (exitCleanupInstalled) return; + exitCleanupInstalled = true; + process.once('exit', cleanupHeldClaims); +}; + +const holderMessage = (holder, { stale = false } = {}) => { + const started = Number.isFinite(holder?.startedAt) ? `, started ${new Date(holder.startedAt).toISOString()}` : ''; + const suffix = stale ? ' The recorded claim is older than the safety ceiling.' : ''; + return `Local accelerator is in use by ${holder?.kind || 'another job'} ${holder?.id || 'unknown'}${started}.${suffix}`; +}; + +const publicHolder = ({ kind, id, pid, startedAt } = {}) => ({ kind, id, pid, startedAt }); + +/** + * Claim the machine-local accelerator. + * + * A zero timeout refuses immediately, suitable for an interactive Generate + * action. Background callers may provide a bounded timeout to wait for a + * currently-running job without silently waiting forever. + */ +export async function claimHeavyLocalJob({ + kind, + id, + timeoutMs = 0, + claimPath = HEAVY_LOCAL_JOB_CLAIM_PATH, + pid = process.pid, + now = () => Date.now(), + wait = sleep, + pidIsAlive = isPidAlive, +} = {}) { + if (typeof kind !== 'string' || !kind.trim() || typeof id !== 'string' || !id.trim()) { + throw new Error('A heavy local job claim requires non-empty kind and id.'); + } + const deadline = now() + Math.max(0, timeoutMs); + const token = randomUUID(); + const ours = { kind, id, pid, startedAt: now(), token }; + await mkdir(dirname(claimPath), { recursive: true }); + + for (;;) { + try { + await writeFile(claimPath, JSON.stringify(ours), { flag: 'wx' }); + heldClaims.set(claimPath, token); + installExitCleanup(); + let released = false; + return { + ok: true, + holder: publicHolder(ours), + // Detached media children outlive a server restart. Transfer the PID + // recorded on disk once one exists so process-exit cleanup does not + // release a claim while that child is still consuming the accelerator. + async handoffTo(childPid) { + if (!Number.isInteger(childPid) || childPid <= 0) return; + const current = await readHolder(claimPath); + if (current?.token !== token) return; + ours.pid = childPid; + await writeFile(claimPath, JSON.stringify(ours)); + heldClaims.delete(claimPath); + }, + async release() { + if (released) return; + released = true; + heldClaims.delete(claimPath); + const current = await readHolder(claimPath); + if (current?.token === token) await unlink(claimPath).catch((err) => { + if (err?.code !== 'ENOENT') throw err; + }); + }, + }; + } catch (err) { + if (err?.code !== 'EEXIST') throw err; + } + + const holder = await readHolder(claimPath); + if (!holder || !pidIsAlive(holder.pid)) { + await unlink(claimPath).catch((err) => { + if (err?.code !== 'ENOENT') throw err; + }); + continue; + } + const stale = now() - holder.startedAt > HEAVY_LOCAL_JOB_STALE_MS; + if (now() >= deadline) { + return { ok: false, holder: publicHolder(holder), stale, message: holderMessage(holder, { stale }), release: async () => {} }; + } + await wait(Math.min(250, Math.max(1, deadline - now()))); + } +} diff --git a/server/lib/heavyJobClaim.test.js b/server/lib/heavyJobClaim.test.js new file mode 100644 index 0000000000..63be885c3a --- /dev/null +++ b/server/lib/heavyJobClaim.test.js @@ -0,0 +1,69 @@ +import { mkdtemp, readFile, rm, writeFile } from 'fs/promises'; +import { existsSync } from 'fs'; +import { tmpdir } from 'os'; +import { join } from 'path'; +import { describe, expect, it, vi } from 'vitest'; +import { claimHeavyLocalJob, HEAVY_LOCAL_JOB_STALE_MS } from './heavyJobClaim.js'; + +const withClaimPath = async (run) => { + const dir = await mkdtemp(join(tmpdir(), 'portos-heavy-job-')); + const claimPath = join(dir, 'claim.json'); + try { + await run(claimPath); + } finally { + await rm(dir, { recursive: true, force: true }); + } +}; + +describe('claimHeavyLocalJob', () => { + it('holds one cross-process claim until its owner releases it', async () => { + await withClaimPath(async (claimPath) => { + const first = await claimHeavyLocalJob({ kind: 'image', id: 'job-a', claimPath }); + expect(first.ok).toBe(true); + expect(first.holder).toMatchObject({ kind: 'image', id: 'job-a', pid: process.pid }); + expect(first.holder).not.toHaveProperty('token'); + + const second = await claimHeavyLocalJob({ kind: 'video', id: 'job-b', claimPath, pidIsAlive: () => true }); + expect(second).toMatchObject({ ok: false, holder: { kind: 'image', id: 'job-a' } }); + + await first.handoffTo(12345); + expect(JSON.parse(await readFile(claimPath, 'utf8'))).toMatchObject({ pid: 12345 }); + await first.release(); + expect(existsSync(claimPath)).toBe(false); + }); + }); + + it('reclaims a lock whose recorded PID is no longer alive', async () => { + await withClaimPath(async (claimPath) => { + await writeFile(claimPath, JSON.stringify({ kind: 'training', id: 'old', pid: 42, startedAt: 1, token: 'old' })); + const claim = await claimHeavyLocalJob({ kind: 'video', id: 'new', claimPath, pidIsAlive: () => false }); + expect(claim).toMatchObject({ ok: true, holder: { kind: 'video', id: 'new' } }); + expect(JSON.parse(await readFile(claimPath, 'utf8'))).toMatchObject({ kind: 'video', id: 'new' }); + await claim.release(); + }); + }); + + it('waits only through its bounded timeout for a background caller', async () => { + await withClaimPath(async (claimPath) => { + await writeFile(claimPath, JSON.stringify({ kind: 'training', id: 'held', pid: 42, startedAt: 0, token: 'old' })); + let now = 0; + const wait = vi.fn(async (ms) => { now += ms; }); + const result = await claimHeavyLocalJob({ + kind: 'video', id: 'waiting', claimPath, timeoutMs: 500, now: () => now, wait, pidIsAlive: () => true, + }); + expect(result).toMatchObject({ ok: false, holder: { kind: 'training', id: 'held' } }); + expect(wait).toHaveBeenCalled(); + }); + }); + + it('reports a live holder that has exceeded the stale-age ceiling without stealing it', async () => { + await withClaimPath(async (claimPath) => { + await writeFile(claimPath, JSON.stringify({ kind: 'training', id: 'held', pid: 42, startedAt: 0, token: 'old' })); + const result = await claimHeavyLocalJob({ + kind: 'video', id: 'next', claimPath, now: () => HEAVY_LOCAL_JOB_STALE_MS + 1, pidIsAlive: () => true, + }); + expect(result.stale).toBe(true); + expect(result.message).toContain('older than the safety ceiling'); + }); + }); +}); diff --git a/server/lib/index.js b/server/lib/index.js index 7b5cc7ba3e..24503586b4 100644 --- a/server/lib/index.js +++ b/server/lib/index.js @@ -108,6 +108,7 @@ export * from './cursor.js'; export * from './grok.js'; export * from './grokVideoClip.js'; export * from './hfToken.js'; +export * from './heavyJobClaim.js'; export * from './hfErrors.js'; export * from './hfCache.js'; export * from './hfDownload.js'; @@ -166,6 +167,7 @@ export * from './schemaVersions.js'; export * from './imageClean.js'; export * from './imageWatermark.js'; export * from './localImageFilename.js'; +export * from './localMemory.js'; export * from './pgFileFacade.js'; export * from './multipart.js'; export * from './safetensors.js'; diff --git a/server/lib/localMemory.js b/server/lib/localMemory.js new file mode 100644 index 0000000000..0b4f81b634 --- /dev/null +++ b/server/lib/localMemory.js @@ -0,0 +1,62 @@ +/** Shared best-effort local-model memory reclamation and headroom reporting. */ + +import { execFile } from 'child_process'; +import { platform, freemem, totalmem } from 'os'; +import { promisify } from 'util'; +import { getLoadedModels as ollamaLoadedModels, unloadModel as ollamaUnload, getBaseUrl as ollamaBaseUrl } from '../services/ollamaManager.js'; +import { getLoadedModels as lmStudioLoadedModels, unloadModel as lmStudioUnload, getBaseUrl as lmStudioBaseUrl } from '../services/lmStudioManager.js'; + +const execFileAsync = promisify(execFile); +const GB = 2 ** 30; +const LOOPBACK_HOSTS = new Set(['localhost', '127.0.0.1', '0.0.0.0', '::1']); + +export function isLocalBackendUrl(url) { + if (!url || !URL.canParse(url)) return false; + const host = new URL(url).hostname.replace(/^\[|\]$/g, ''); + return LOOPBACK_HOSTS.has(host) || host.startsWith('127.'); +} + +export async function unloadResidentModels() { + const unloaded = []; + if (isLocalBackendUrl(ollamaBaseUrl())) { + const models = await ollamaLoadedModels().catch(() => []); + for (const model of models) { + const name = model?.name || model?.id; + const result = name ? await ollamaUnload(name).catch(() => null) : null; + if (result?.unloaded) unloaded.push(`ollama:${name}`); + } + } + if (isLocalBackendUrl(lmStudioBaseUrl())) { + const models = await lmStudioLoadedModels(true).catch(() => []); + for (const model of models) { + const result = model?.id ? await lmStudioUnload(model.id).catch(() => null) : null; + if (result?.success) unloaded.push(`lmstudio:${model.id}`); + } + } + return unloaded; +} + +const parsePageSize = (out) => Number(out.match(/page size of (\d+) bytes/i)?.[1] || 4096); + +async function darwinAvailableGb() { + const { stdout } = await execFileAsync('vm_stat'); + const pageSize = parsePageSize(stdout); + const pages = (label) => Number(stdout.match(new RegExp(`${label}:\\s+(\\d+)\\.`))?.[1] || 0); + const available = pages('Pages free') + pages('Pages inactive') + pages('Pages speculative') + pages('Pages purgeable'); + return available ? (available * pageSize) / GB : null; +} + +export async function getAvailableMemoryGb() { + if (platform() === 'darwin') { + const available = await darwinAvailableGb().catch(() => null); + if (Number.isFinite(available) && available > 0) return available; + } + return freemem() / GB; +} + +export async function prepareLocalMemory() { + const unloaded = await unloadResidentModels().catch(() => []); + const availableGb = await getAvailableMemoryGb().catch(() => 0); + const totalGb = totalmem() / GB; + return { unloaded, availableGb, totalGb, budgetGb: Math.min(totalGb, availableGb) }; +} diff --git a/server/services/imageGen/local.js b/server/services/imageGen/local.js index a61225353c..70b5018bf8 100644 --- a/server/services/imageGen/local.js +++ b/server/services/imageGen/local.js @@ -29,6 +29,8 @@ import { hfChildEnv } from '../../lib/hfToken.js'; import { extractGatedRepo, isGatedRepoError } from '../../lib/hfErrors.js'; import { killWithEscalation } from '../../lib/killWithEscalation.js'; import { createLineReader } from '../../lib/streamLines.js'; +import { claimHeavyLocalJob } from '../../lib/heavyJobClaim.js'; +import { prepareLocalMemory } from '../../lib/localMemory.js'; import { IMAGE_GEN_MODE, LOCAL_IMAGEGEN_DEFAULT_MODEL } from './modes.js'; import { computePixelDelta } from './regen.js'; import { parseByteProgress, formatDownloadMessage } from '../videoGen/generateVideoHelpers.js'; @@ -541,6 +543,16 @@ export async function generateImage({ pythonPath, prompt = '', negativePrompt = const stepwiseDir = await mkdtemp(join(tmpdir(), 'portos-stepwise-')); const { bin, args } = buildArgs({ pythonPath, model, prompt, negativePrompt, width: Number(width), height: Number(height), steps: actualSteps, guidance: actualGuidance, seed: actualSeed, quantize, outputPath, loraPaths: validLoras, loraScales, stepwiseDir, initImagePath: validInitImagePath, initImageStrength: validInitImageStrength, referenceImagePaths: validReferenceImagePaths, referenceImageStrengths: validReferenceImageStrengths }); + const heavyClaim = await claimHeavyLocalJob({ kind: 'local image generation', id: jobId }); + if (!heavyClaim.ok) { + jobs.delete(jobId); + await rm(stepwiseDir, { recursive: true, force: true }); + throw new ServerError(heavyClaim.message, { status: 409, code: 'HEAVY_LOCAL_JOB_BUSY', context: { holder: heavyClaim.holder } }); + } + const releaseHeavyClaim = () => heavyClaim.release() + .catch((err) => console.error(`❌ Image generation claim release [${jobId.slice(0, 8)}]: ${err.message}`)); + const memoryReport = await prepareLocalMemory(); + if (memoryReport.unloaded.length) console.log(`🧹 Image generation [${jobId.slice(0, 8)}] freed ${memoryReport.unloaded.length} resident model(s)`); console.log(`🎨 Generating image [${jobId.slice(0, 8)}] local: ${modelId} ${width}x${height} steps=${actualSteps}`); imageGenEvents.emit('started', { generationId: jobId, totalSteps: actualSteps }); @@ -548,6 +560,7 @@ export async function generateImage({ pythonPath, prompt = '', negativePrompt = const proc = spawn(bin, args, { env: await hfChildEnv(), stdio: ['ignore', 'pipe', 'pipe'] }); activeProcess = proc; + await heavyClaim.handoffTo?.(proc.pid); // Spawn ENOENT (missing/non-executable pythonPath) fires BOTH 'error' and // 'close' on Node — without this guard, a typo'd pythonPath emits two // 'failed' events to imageGenEvents and two SSE error frames to the @@ -566,6 +579,7 @@ export async function generateImage({ pythonPath, prompt = '', negativePrompt = imageGenEvents.emit('failed', { mode: IMAGE_GEN_MODE.LOCAL, generationId: jobId, error: reason }); activeProcess = null; activeJob = null; + void releaseHeavyClaim(); rm(stepwiseDir, { recursive: true, force: true }).catch(() => {}); closeJobAfterDelay(jobs, jobId); }); @@ -741,6 +755,7 @@ export async function generateImage({ pythonPath, prompt = '', negativePrompt = stdoutReader.flush(); activeProcess = null; activeJob = null; + void releaseHeavyClaim(); if (watcher) { try { watcher.close(); } catch { /* ignore */ } } rm(stepwiseDir, { recursive: true, force: true }).catch(() => {}); if (code !== 0) { diff --git a/server/services/imageTo3d/models.js b/server/services/imageTo3d/models.js index 0fffceea49..601ce8e95b 100644 --- a/server/services/imageTo3d/models.js +++ b/server/services/imageTo3d/models.js @@ -20,6 +20,8 @@ import { join } from 'node:path'; import { rm } from 'node:fs/promises'; import { ServerError } from '../../lib/errorHandler.js'; import { PATHS, resolveGalleryImage, ensureDir } from '../../lib/fileUtils.js'; +import { claimHeavyLocalJob } from '../../lib/heavyJobClaim.js'; +import { prepareLocalMemory } from '../../lib/localMemory.js'; import { slugifyForFilename } from '../../lib/civitai.js'; import { detectHostCapabilities, resolveTarget, DEFAULT_IMAGE_TO_3D_TARGET } from './targets.js'; import { getTargetAdapter } from './adapters.js'; @@ -113,8 +115,15 @@ async function failGeneration(id, operationId, error) { async function executeRender({ id, operationId, adapter, sourcePath, caps }) { const outputPath = assetDiskPath(id); let lastPersistedPercent = -1; + let heavyClaim = null; try { await ensureDir(join(PATHS.imageTo3d, id)); + heavyClaim = await claimHeavyLocalJob({ kind: 'image-to-3D generation', id: operationId }); + if (!heavyClaim.ok) { + throw new ServerError(heavyClaim.message, { status: 409, code: 'HEAVY_LOCAL_JOB_BUSY', context: { holder: heavyClaim.holder } }); + } + const memoryReport = await prepareLocalMemory(); + if (memoryReport.unloaded.length) console.log(`🧹 Image-to-3D render freed ${memoryReport.unloaded.length} resident model(s)`); // Resolve this target's own credential/env needs via its adapter (e.g. // TRELLIS.2 resolves the stored Hugging Face token) — omitted for a target // with nothing to resolve. Resolving HERE (an async caller) keeps `run` @@ -179,6 +188,7 @@ async function executeRender({ id, operationId, adapter, sourcePath, caps }) { console.error(`❌ Image-to-3D render failed for ${id}: ${cleanError(error)}`); await failGeneration(id, operationId, error); } finally { + await heavyClaim?.release().catch((err) => console.error(`❌ Image-to-3D claim release failed: ${err.message}`)); activeRenders.delete(operationId); activeOperations.delete(operationId); // If the record was deleted while the render ran, the completion/failure writes diff --git a/server/services/imageTo3d/models.test.js b/server/services/imageTo3d/models.test.js index fb3dd17560..3f12771d04 100644 --- a/server/services/imageTo3d/models.test.js +++ b/server/services/imageTo3d/models.test.js @@ -40,6 +40,14 @@ vi.mock('../../lib/hfToken.js', () => ({ hfChildEnv: vi.fn(async () => ({ HF_TOKEN: 'hf_from_store', HUGGINGFACE_HUB_TOKEN: 'hf_from_store' })), })); +const { claimRelease } = vi.hoisted(() => ({ claimRelease: vi.fn(async () => {}) })); +vi.mock('../../lib/heavyJobClaim.js', () => ({ + claimHeavyLocalJob: vi.fn(async () => ({ ok: true, holder: {}, release: claimRelease })), +})); +vi.mock('../../lib/localMemory.js', () => ({ + prepareLocalMemory: vi.fn(async () => ({ unloaded: [], availableGb: 64, totalGb: 64, budgetGb: 64 })), +})); + vi.mock('./db.js', () => ({ listModels: vi.fn(), getModel: vi.fn(), @@ -53,6 +61,7 @@ import { rm } from 'node:fs/promises'; import { ensureDir } from '../../lib/fileUtils.js'; import { resolveTarget } from './targets.js'; import { isTrellis2Installed, runTrellis2Generate } from './trellis2.js'; +import { claimHeavyLocalJob } from '../../lib/heavyJobClaim.js'; import * as store from './db.js'; import { createModel, startGeneration, getModelAsset, recoverInterruptedModels, deleteModel, @@ -72,6 +81,7 @@ const draftRecord = () => ({ describe('image-to-3D model orchestration', () => { beforeEach(() => { vi.clearAllMocks(); + claimHeavyLocalJob.mockResolvedValue({ ok: true, holder: {}, release: claimRelease }); isTrellis2Installed.mockReturnValue(true); resolveTarget.mockImplementation((id) => ( id === 'trellis2' @@ -139,6 +149,8 @@ describe('image-to-3D model orchestration', () => { // #3032: the resolved HF token (settings-stored included) rides into the child, // merged over process.env — without it, gated DINOv3/RMBG-2.0 pulls 401. expect(generateArgs.env).toMatchObject({ HF_TOKEN: 'hf_from_store', HUGGINGFACE_HUB_TOKEN: 'hf_from_store' }); + expect(claimHeavyLocalJob).toHaveBeenCalledWith(expect.objectContaining({ kind: 'image-to-3D generation' })); + expect(claimRelease).toHaveBeenCalled(); expect(posixPath(current.assetPath)).toBe('/data/image-to-3d/image3d-example/model.glb'); expect(current.generationOperationId).toBeNull(); expect(current.runs.at(-1)).toMatchObject({ status: 'completed', percent: 100 }); diff --git a/server/services/loraTraining/index.js b/server/services/loraTraining/index.js index fd79148220..45cf8526ac 100644 --- a/server/services/loraTraining/index.js +++ b/server/services/loraTraining/index.js @@ -43,6 +43,7 @@ import { import { makeTrainingLineHandler } from './progress.js'; import { makeStallDetector } from './stallDetector.js'; import { prepareMemoryForTraining, TRAINING_MIN_HEADROOM_GB } from './memoryPrep.js'; +import { claimHeavyLocalJob } from '../../lib/heavyJobClaim.js'; import { classifyTrainingFailure } from './failure.js'; import { buildTrainedSidecar, trainedLoraFilename } from './sidecar.js'; import { validateDatasetReady } from './dataset.js'; @@ -440,6 +441,7 @@ export async function runTraining({ jobId, runId, pythonPath = null, resumeCheck if (!run) return fail(`run record missing: ${runId}`); const settings = await getSettings(); const dir = runDir(runId); + let heavyClaim = null; // Terminal failure BEFORE the child spawns: flip the run record to failed // AND release the dataset's `training` status, then emit the failed event. @@ -447,6 +449,7 @@ export async function runTraining({ jobId, runId, pythonPath = null, resumeCheck // stuck `running` (lingering until the next boot reconcile) or the dataset // stuck on its `training` chip. const failBeforeSpawn = async (message) => { + await heavyClaim?.release().catch((err) => console.error(`❌ training [${shortId(jobId)}] claim release failed: ${err.message}`)); await runsDb.updateRun(runId, { status: 'failed', error: message, completedAt: new Date().toISOString(), }).catch(() => {}); @@ -454,6 +457,9 @@ export async function runTraining({ jobId, runId, pythonPath = null, resumeCheck return fail(message); }; + heavyClaim = await claimHeavyLocalJob({ kind: 'LoRA training', id: jobId }); + if (!heavyClaim.ok) return failBeforeSpawn(heavyClaim.message); + // Boot re-attach (#1332): the queue re-enqueues a run whose detached trainer // SURVIVED a hard server restart with `reattach: true`. The child reparented // to init and is still training (or finished during the downtime), so we tail @@ -634,7 +640,12 @@ export async function runTraining({ jobId, runId, pythonPath = null, resumeCheck // and the queue worker awaits the 'close' event for the run lifecycle. No // `cleanup` — those logs are the only copy of raw trainer stdout/stderr, kept // in the run dir for post-mortem and removed when the run dir is deleted. - const proc = await spawnDetached(bin, args, { env: childEnv, controlDir: join(dir, '.detached') }); + let proc; + try { + proc = await spawnDetached(bin, args, { env: childEnv, controlDir: join(dir, '.detached') }); + } catch (err) { + return failBeforeSpawn(`trainer spawn failed: ${err.message}`); + } wireProcLifecycle(proc); // Wire a freshly-spawned OR re-attached (#1332) trainer handle into the full @@ -646,6 +657,7 @@ export async function runTraining({ jobId, runId, pythonPath = null, resumeCheck function wireProcLifecycle(proc, { isReattach = false } = {}) { activeProcess = proc; activeJobId = jobId; + void heavyClaim?.handoffTo?.(proc.pid)?.catch((err) => console.error(`❌ training [${shortId(jobId)}] claim handoff failed: ${err.message}`)); // Keep the Mac awake for the duration of the training child (idle + system // sleep held off) — but DELIBERATELY let the *display* sleep. An active @@ -836,7 +848,11 @@ export async function runTraining({ jobId, runId, pythonPath = null, resumeCheck // finalize reads the run record — the collapse guard and previewImageUrl // both read run.artifacts. Async finalize wrapped so a rejection can't // escape the event handler (unhandled rejection kills the process on Node ≥15). - Promise.resolve(flushProgress()) + // The trainer has exited, so release before finalization can enqueue an + // automatic checkpoint resume; that successor must acquire a fresh claim. + Promise.resolve(heavyClaim?.release()) + .catch((err) => console.error(`❌ training [${shortId(jobId)}] claim release failed: ${err.message}`)) + .then(() => flushProgress()) .then(() => finalizeTraining({ jobId, runId, code, signal, state: getState(), stallKilled })) .then((resumed) => { // Wake the display now the run is over so the user sees the result — diff --git a/server/services/loraTraining/memoryPrep.js b/server/services/loraTraining/memoryPrep.js index f7951bdc0e..1d50930418 100644 --- a/server/services/loraTraining/memoryPrep.js +++ b/server/services/loraTraining/memoryPrep.js @@ -1,144 +1,11 @@ -/** - * Memory preparation for LoRA training runs. - * - * Training shares the Apple-Silicon unified-memory pool with every other - * resident model server (ollama, LM Studio) and with PortOS itself. A run that - * oversubscribes the pool swap-thrashes — ~21 GB of swap was live during the - * GPU watchdog-timeout reboots documented in - * docs/research/2026-06-13-mflux-training-watchdog-panic.md. Before spawning a - * trainer we (1) unload resident LLMs to reclaim their memory and (2) measure - * the real available headroom so the caller can size the run config to what's - * actually free and refuse to start (rather than crash mid-run) when the pool - * is too tight. - * - * Everything here is best-effort and never throws — a failed unload or an - * unreadable vm_stat must not block training; it just yields a more - * conservative (smaller) memory budget. - */ +/** Training-specific memory policy built on the shared local-memory preflight. */ -import { execFile } from 'child_process'; -import { promisify } from 'util'; -import { platform, freemem, totalmem } from 'os'; -import { getLoadedModels as ollamaLoadedModels, unloadModel as ollamaUnload, getBaseUrl as ollamaBaseUrl } from '../ollamaManager.js'; -import { getLoadedModels as lmStudioLoadedModels, unloadModel as lmStudioUnload, getBaseUrl as lmStudioBaseUrl } from '../lmStudioManager.js'; +import { prepareLocalMemory } from '../../lib/localMemory.js'; -const execFileAsync = promisify(execFile); -const GB = 2 ** 30; +export { unloadResidentModels, getAvailableMemoryGb, isLocalBackendUrl } from '../../lib/localMemory.js'; -// Headroom floor: below this even a 4-bit 4B run risks swap-thrash, so refuse -// to start rather than reboot the box. Tuned for the smallest supported run; -// larger variants are protected by the budget-derived quantize/low_ram tiers. +// Below this floor even a small 4-bit training run risks swap-thrash. Larger +// run sizing remains training-specific and still consumes `budgetGb` below. export const TRAINING_MIN_HEADROOM_GB = 24; -// Loopback hosts, plus `0.0.0.0` — a backend URL of `0.0.0.0` targets the local -// machine (the all-interfaces bind, which clients reach via loopback), so it is -// "local" for the purpose of deciding whether an unload frees local memory. -const LOOPBACK_HOSTS = new Set(['localhost', '127.0.0.1', '0.0.0.0', '::1']); - -/** - * True only when `url` points at this machine (loopback, or the `0.0.0.0` - * local-bind sentinel). Unloading a - * model frees memory on the box the backend RUNS on, not the box that issued - * the request — and PortOS supports pointing `OLLAMA_URL` / `LM_STUDIO_URL` at a - * remote LAN peer (a common federated-machines setup). Evicting a *remote* - * backend would free no local unified memory and would destroy another box's - * loaded model for nothing, so we only unload when the backend is local. - * Unparseable or non-loopback → treated as remote (skip the unload). All of - * 127.0.0.0/8 is loopback, so any `127.*` host counts. - */ -export function isLocalBackendUrl(url) { - if (!url || !URL.canParse(url)) return false; - // URL.hostname keeps the [...] brackets on IPv6 literals; strip them so the - // `::1` loopback compares against LOOPBACK_HOSTS. - const host = new URL(url).hostname.replace(/^\[|\]$/g, ''); - return LOOPBACK_HOSTS.has(host) || host.startsWith('127.'); -} - -/** - * Best-effort: unload every model currently resident in ollama and LM Studio so - * its unified memory returns to the pool before a training run. Each unload is - * independent — failures (server down, model already expired) are swallowed. - * Skips a backend whose configured URL is NOT loopback-local: a remote backend - * doesn't share this machine's memory, so evicting it would free nothing here - * and needlessly drop a peer's loaded model. Returns the freed model labels for - * logging. Never throws. - */ -export async function unloadResidentModels() { - const unloaded = []; - - if (isLocalBackendUrl(ollamaBaseUrl())) { - const ollamaLoaded = await ollamaLoadedModels().catch(() => []); - for (const m of ollamaLoaded) { - const name = m?.name || m?.id; - if (!name) continue; - const res = await ollamaUnload(name).catch(() => null); - if (res?.unloaded) unloaded.push(`ollama:${name}`); - } - } - - if (isLocalBackendUrl(lmStudioBaseUrl())) { - const lmLoaded = await lmStudioLoadedModels(true).catch(() => []); - for (const m of lmLoaded) { - if (!m?.id) continue; - const res = await lmStudioUnload(m.id).catch(() => null); - if (res?.success) unloaded.push(`lmstudio:${m.id}`); - } - } - - return unloaded; -} - -const parsePageSize = (out) => { - const m = out.match(/page size of (\d+) bytes/i); // case-insensitive for robustness - return m ? Number(m[1]) : 4096; -}; - -/** - * Parse `vm_stat` for the memory macOS can actually hand to a new process — - * free + inactive + speculative + purgeable pages (the same buckets Activity - * Monitor reclaims under pressure). Node's freemem() counts only truly-free - * pages and so wildly understates available unified memory on macOS, which - * keeps most RAM as reclaimable cache. Returns GB, or null if unparseable. - */ -async function darwinAvailableGb() { - const { stdout } = await execFileAsync('vm_stat'); - const pageSize = parsePageSize(stdout); - const pages = (label) => { - const m = stdout.match(new RegExp(`${label}:\\s+(\\d+)\\.`)); - return m ? Number(m[1]) : 0; - }; - const available = pages('Pages free') + pages('Pages inactive') - + pages('Pages speculative') + pages('Pages purgeable'); - if (!available) return null; - return (available * pageSize) / GB; -} - -/** - * Memory the OS can realistically give a training run right now, in GB. Uses - * vm_stat on darwin (unified-memory-aware), falling back to freemem() on other - * platforms or any failure. Never throws. - */ -export async function getAvailableMemoryGb() { - if (platform() === 'darwin') { - const v = await darwinAvailableGb().catch(() => null); - if (Number.isFinite(v) && v > 0) return v; - } - return freemem() / GB; -} - -/** - * Reclaim memory and report the budget for sizing/gating a training run: - * - unloaded: labels of resident models freed - * - availableGb: memory free after unloading - * - totalGb: physical RAM - * - budgetGb: what training may use (available, clamped to physical) — feed - * this to deriveMfluxMemoryConfig so the quantize/low_ram tier reflects - * real headroom, not raw RAM. Never throws. - */ -export async function prepareMemoryForTraining() { - const unloaded = await unloadResidentModels().catch(() => []); - const availableGb = await getAvailableMemoryGb().catch(() => 0); - const totalGb = totalmem() / GB; - const budgetGb = Math.min(totalGb, availableGb); - return { unloaded, availableGb, totalGb, budgetGb }; -} +export const prepareMemoryForTraining = prepareLocalMemory; diff --git a/server/services/loraTraining/memoryPrep.test.js b/server/services/loraTraining/memoryPrep.test.js index 1767c1de24..6e6c9e6c18 100644 --- a/server/services/loraTraining/memoryPrep.test.js +++ b/server/services/loraTraining/memoryPrep.test.js @@ -15,12 +15,12 @@ const h = vi.hoisted(() => ({ vmStatThrows: false, })); -vi.mock('../ollamaManager.js', () => ({ +vi.mock('../../services/ollamaManager.js', () => ({ getLoadedModels: vi.fn(async () => h.ollamaLoaded), unloadModel: (...a) => h.ollamaUnload(...a), getBaseUrl: () => h.ollamaUrl, })); -vi.mock('../lmStudioManager.js', () => ({ +vi.mock('../../services/lmStudioManager.js', () => ({ getLoadedModels: vi.fn(async () => h.lmLoaded), unloadModel: (...a) => h.lmUnload(...a), getBaseUrl: () => h.lmUrl, diff --git a/server/services/videoGen/local.js b/server/services/videoGen/local.js index 94cf1aa99b..100bce8a54 100644 --- a/server/services/videoGen/local.js +++ b/server/services/videoGen/local.js @@ -21,6 +21,8 @@ import { ensureDir, PATHS, UUID_RE } from '../../lib/fileUtils.js'; import { spawnDetached } from '../../lib/detachedSpawn.js'; import { killWithEscalation } from '../../lib/killWithEscalation.js'; import { createLineReader } from '../../lib/streamLines.js'; +import { claimHeavyLocalJob } from '../../lib/heavyJobClaim.js'; +import { prepareLocalMemory } from '../../lib/localMemory.js'; import { ServerError } from '../../lib/errorHandler.js'; import { videoLoraLayoutIssue } from '../../lib/safetensors.js'; import { videoGenEvents } from './events.js'; @@ -1508,6 +1510,22 @@ export async function generateVideo({ pythonPath, prompt, negativePrompt = '', m throw err; } + const heavyClaim = await claimHeavyLocalJob({ kind: 'local video generation', id: jobId }); + if (!heavyClaim.ok) { + jobs.delete(jobId); + if (resizedSrcTempPath) await unlink(resizedSrcTempPath).catch(() => {}); + if (resizedLastTempPath) await unlink(resizedLastTempPath).catch(() => {}); + await Promise.all(resizedKeyframeTempPaths.map((path) => unlink(path).catch(() => {}))); + await Promise.all(icReferenceTempPaths.map((path) => unlink(path).catch(() => {}))); + if (uploadedTempPath) await unlink(uploadedTempPath).catch(() => {}); + await Promise.all(uploadedTempPaths.map((path) => unlink(path).catch(() => {}))); + throw new ServerError(heavyClaim.message, { status: 409, code: 'HEAVY_LOCAL_JOB_BUSY', context: { holder: heavyClaim.holder } }); + } + const releaseHeavyClaim = () => heavyClaim.release() + .catch((err) => console.error(`❌ Video generation claim release [${jobId.slice(0, 8)}]: ${err.message}`)); + const memoryReport = await prepareLocalMemory(); + if (memoryReport.unloaded.length) console.log(`🧹 Video generation [${jobId.slice(0, 8)}] freed ${memoryReport.unloaded.length} resident model(s)`); + // History-calibrated wall-clock estimate (#3801). `null` when this install // has never measured a render on this model — an explicit "no estimate" // sentinel the UI must render as "unknown", never as 0 or a guess. Stamped @@ -1564,13 +1582,20 @@ export async function generateVideo({ pythonPath, prompt, negativePrompt = '', m // still `proc.kill()` it directly by PID on cancel / watchdog. `cleanup: true` // lets the helper drop that scratch dir on every terminal path (close/error) // so it can't accumulate under data/videos. - const proc = await spawnDetached(bin, args, { - env: childEnv, - controlDir: join(PATHS.videos, '.detached', jobId), - cleanup: true, - killProcessGroup: runtimeNeedsProcessGroupKill(model.runtime), - }); + let proc; + try { + proc = await spawnDetached(bin, args, { + env: childEnv, + controlDir: join(PATHS.videos, '.detached', jobId), + cleanup: true, + killProcessGroup: runtimeNeedsProcessGroupKill(model.runtime), + }); + } catch (err) { + await releaseHeavyClaim(); + throw err; + } activeProcess = proc; + await heavyClaim.handoffTo?.(proc.pid); // Panel-side completion watchdog. Armed once we see the render's completion // marker on stdout; SIGKILLs the child if it hasn't exited after the grace @@ -1680,6 +1705,7 @@ export async function generateVideo({ pythonPath, prompt, negativePrompt = '', m broadcastSse(job, { type: 'error', error: reason }); videoGenEvents.emit('failed', { generationId: jobId, error: reason }); activeProcess = null; + void releaseHeavyClaim(); // Spawn failed, so proc.on('close') will never fire — clean up every // temp file we own here, including the multipart upload, otherwise // ENOENT/permission errors leak files in os.tmpdir(). @@ -1765,6 +1791,10 @@ export async function generateVideo({ pythonPath, prompt, negativePrompt = '', m clearCompletionWatchdog(); clearIdleStallTimer(); activeProcess = null; + // The child has exited, so its accelerator allocation is gone. Release + // before emitting the terminal completion event: an extend chain starts its + // next child from that event and must be able to acquire the machine claim. + await releaseHeavyClaim(); // Wrap the whole teardown so a throw from finalizeGeneratedVideo (history // save, thumbnail, file move) can't leak as an unhandled rejection — on // Node ≥15 that kills the process AND strands the media job `running` with diff --git a/server/services/videoGen/local.test.js b/server/services/videoGen/local.test.js index 66b9b22f41..43afebcac9 100644 --- a/server/services/videoGen/local.test.js +++ b/server/services/videoGen/local.test.js @@ -10,6 +10,14 @@ import { basename, join } from 'path'; import { tmpdir, totalmem } from 'os'; import { randomUUID } from 'crypto'; +const { heavyClaimRelease } = vi.hoisted(() => ({ heavyClaimRelease: vi.fn(async () => {}) })); +vi.mock('../../lib/heavyJobClaim.js', () => ({ + claimHeavyLocalJob: vi.fn(async () => ({ ok: true, holder: {}, release: heavyClaimRelease })), +})); +vi.mock('../../lib/localMemory.js', () => ({ + prepareLocalMemory: vi.fn(async () => ({ unloaded: [], availableGb: 64, totalGb: 64, budgetGb: 64 })), +})); + // ─── dep mocks (must be declared before the module import) ─────────────────── const MOCK_PATHS = { From 9675d96c018b1864e1de9745cc30f3498c9ec240 Mon Sep 17 00:00:00 2001 From: "[._.]/ Adam Eivy" Date: Sat, 15 Aug 2026 07:24:50 -0700 Subject: [PATCH 2/3] fix: mock heavyJobClaim in image-to-3D dispatch test The test mocks fileUtils.js down to PATHS.imageTo3d alone, but models.js also imports claimHeavyLocalJob from heavyJobClaim.js, which reads PATHS.data at module load time. With PATHS.data undefined, the real heavyJobClaim.js threw at import and failed the whole test file in CI. models.test.js already mocks heavyJobClaim.js the same way; this test file needed the same mock. --- server/services/imageTo3d/models.genericDispatch.test.js | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/server/services/imageTo3d/models.genericDispatch.test.js b/server/services/imageTo3d/models.genericDispatch.test.js index 3b7e7b11da..050a4a54a4 100644 --- a/server/services/imageTo3d/models.genericDispatch.test.js +++ b/server/services/imageTo3d/models.genericDispatch.test.js @@ -30,6 +30,14 @@ vi.mock('../../lib/hfToken.js', () => ({ hfChildEnv: vi.fn(async () => ({ HF_TOKEN: 'hf_test' })), })); +// models.js also claims the machine-wide heavy-accelerator lock before +// rendering (see heavyJobClaim.js). Mock it the same way models.test.js does — +// this suite mocks fileUtils.js down to `imageTo3d` alone, and the real +// heavyJobClaim.js needs PATHS.data at import time. +vi.mock('../../lib/heavyJobClaim.js', () => ({ + claimHeavyLocalJob: vi.fn(async () => ({ ok: true, holder: {}, release: vi.fn(() => Promise.resolve()) })), +})); + vi.mock('./db.js', () => ({ listModels: vi.fn(), getModel: vi.fn(), From 398d5be909d3f4d3f0212ff87231e91b25784fb6 Mon Sep 17 00:00:00 2001 From: "[._.]/ Adam Eivy" Date: Sat, 15 Aug 2026 07:33:43 -0700 Subject: [PATCH 3/3] address review (claude): fix boot-reattach deadlocking on its own heavy-job claim MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit runTraining() acquired a fresh machine-wide accelerator claim before checking reattach:true. On a server restart that re-attaches to a detached trainer which survived the crash, that trainer's own PID already holds the claim (handed off pre-restart), so the fresh claimHeavyLocalJob() call saw it as a competing job and refused — failing every restart-survived training run outright, even though it was actively still training. Add adoptHeavyLocalJob() to recognize and adopt an on-disk claim already recorded for this exact kind/id/pid instead of contending for a new one; fall back to a fresh claim only when no matching claim survived (e.g. a run predating this feature). --- server/lib/heavyJobClaim.js | 35 +++++++++++++++++ server/lib/heavyJobClaim.test.js | 55 ++++++++++++++++++++++++++- server/services/loraTraining/index.js | 22 +++++++++-- 3 files changed, 107 insertions(+), 5 deletions(-) diff --git a/server/lib/heavyJobClaim.js b/server/lib/heavyJobClaim.js index 7a2a2ca6d1..8edd0152a2 100644 --- a/server/lib/heavyJobClaim.js +++ b/server/lib/heavyJobClaim.js @@ -143,3 +143,38 @@ export async function claimHeavyLocalJob({ await wait(Math.min(250, Math.max(1, deadline - now()))); } } + +/** + * Adopt a claim already recorded on disk, for a process that did not + * acquire it — a restarted server re-attaching to a detached child that + * survived the crash (#1332). The surviving child already holds the + * machine-wide accelerator claim (transferred to it via `handoffTo` before + * the restart); calling `claimHeavyLocalJob` again here would see that live + * claim as a COMPETING job and refuse it, failing every restart-survived run + * outright. This never contends for the lock — it only recognizes a claim + * that already names this exact job and PID, so it's safe to call blind. + * + * Returns `null` when no on-disk claim matches `kind`/`id`/`pid` (e.g. an + * in-flight run that predates this claim file, or a genuinely different + * holder), so the caller can fall back to `claimHeavyLocalJob`. + */ +export async function adoptHeavyLocalJob({ + kind, id, pid, claimPath = HEAVY_LOCAL_JOB_CLAIM_PATH, +} = {}) { + const holder = await readHolder(claimPath); + if (holder?.kind !== kind || holder.id !== id || holder.pid !== pid || typeof holder.token !== 'string') return null; + let released = false; + return { + ok: true, + holder: publicHolder(holder), + async handoffTo() {}, + async release() { + if (released) return; + released = true; + const current = await readHolder(claimPath); + if (current?.token === holder.token) await unlink(claimPath).catch((err) => { + if (err?.code !== 'ENOENT') throw err; + }); + }, + }; +} diff --git a/server/lib/heavyJobClaim.test.js b/server/lib/heavyJobClaim.test.js index 63be885c3a..afec6f1c72 100644 --- a/server/lib/heavyJobClaim.test.js +++ b/server/lib/heavyJobClaim.test.js @@ -3,7 +3,7 @@ import { existsSync } from 'fs'; import { tmpdir } from 'os'; import { join } from 'path'; import { describe, expect, it, vi } from 'vitest'; -import { claimHeavyLocalJob, HEAVY_LOCAL_JOB_STALE_MS } from './heavyJobClaim.js'; +import { adoptHeavyLocalJob, claimHeavyLocalJob, HEAVY_LOCAL_JOB_STALE_MS } from './heavyJobClaim.js'; const withClaimPath = async (run) => { const dir = await mkdtemp(join(tmpdir(), 'portos-heavy-job-')); @@ -67,3 +67,56 @@ describe('claimHeavyLocalJob', () => { }); }); }); + +describe('adoptHeavyLocalJob', () => { + // #1332 boot re-attach: a restarted server re-attaches to a detached trainer + // that survived the crash. That trainer's PID already holds the machine-wide + // claim (handed off to it pre-restart) — the new process must recognize and + // adopt that claim rather than contending for a fresh one, which would see + // its own survivor as a competing job and refuse it. + it('adopts an on-disk claim already recorded for this exact kind/id/pid', async () => { + await withClaimPath(async (claimPath) => { + await writeFile(claimPath, JSON.stringify({ + kind: 'LoRA training', id: 'run-a', pid: 4242, startedAt: 0, token: 'tok-a', + })); + const adopted = await adoptHeavyLocalJob({ kind: 'LoRA training', id: 'run-a', pid: 4242, claimPath }); + expect(adopted).toMatchObject({ ok: true, holder: { kind: 'LoRA training', id: 'run-a', pid: 4242 } }); + + await adopted.release(); + expect(existsSync(claimPath)).toBe(false); + }); + }); + + it('returns null (does not steal) when no claim exists, or the on-disk claim names a different job', async () => { + await withClaimPath(async (claimPath) => { + expect(await adoptHeavyLocalJob({ kind: 'LoRA training', id: 'run-a', pid: 4242, claimPath })).toBeNull(); + + await writeFile(claimPath, JSON.stringify({ + kind: 'LoRA training', id: 'run-a', pid: 4242, startedAt: 0, token: 'tok-a', + })); + // Wrong id, wrong pid, and wrong kind all fail to match — and the claim + // file is left untouched (still readable, still holding the same token). + expect(await adoptHeavyLocalJob({ kind: 'LoRA training', id: 'run-b', pid: 4242, claimPath })).toBeNull(); + expect(await adoptHeavyLocalJob({ kind: 'LoRA training', id: 'run-a', pid: 9999, claimPath })).toBeNull(); + expect(await adoptHeavyLocalJob({ kind: 'video', id: 'run-a', pid: 4242, claimPath })).toBeNull(); + expect(JSON.parse(await readFile(claimPath, 'utf8'))).toMatchObject({ token: 'tok-a' }); + }); + }); + + it('release() only removes the claim file if its token still matches', async () => { + await withClaimPath(async (claimPath) => { + await writeFile(claimPath, JSON.stringify({ + kind: 'LoRA training', id: 'run-a', pid: 4242, startedAt: 0, token: 'tok-a', + })); + const adopted = await adoptHeavyLocalJob({ kind: 'LoRA training', id: 'run-a', pid: 4242, claimPath }); + + // Someone else already claimed the path under a different token by the + // time release() runs — it must not remove that unrelated claim. + await writeFile(claimPath, JSON.stringify({ + kind: 'video', id: 'run-c', pid: 1, startedAt: 0, token: 'tok-c', + })); + await adopted.release(); + expect(JSON.parse(await readFile(claimPath, 'utf8'))).toMatchObject({ token: 'tok-c' }); + }); + }); +}); diff --git a/server/services/loraTraining/index.js b/server/services/loraTraining/index.js index 45cf8526ac..858425a2fb 100644 --- a/server/services/loraTraining/index.js +++ b/server/services/loraTraining/index.js @@ -43,7 +43,7 @@ import { import { makeTrainingLineHandler } from './progress.js'; import { makeStallDetector } from './stallDetector.js'; import { prepareMemoryForTraining, TRAINING_MIN_HEADROOM_GB } from './memoryPrep.js'; -import { claimHeavyLocalJob } from '../../lib/heavyJobClaim.js'; +import { claimHeavyLocalJob, adoptHeavyLocalJob } from '../../lib/heavyJobClaim.js'; import { classifyTrainingFailure } from './failure.js'; import { buildTrainedSidecar, trainedLoraFilename } from './sidecar.js'; import { validateDatasetReady } from './dataset.js'; @@ -457,9 +457,6 @@ export async function runTraining({ jobId, runId, pythonPath = null, resumeCheck return fail(message); }; - heavyClaim = await claimHeavyLocalJob({ kind: 'LoRA training', id: jobId }); - if (!heavyClaim.ok) return failBeforeSpawn(heavyClaim.message); - // Boot re-attach (#1332): the queue re-enqueues a run whose detached trainer // SURVIVED a hard server restart with `reattach: true`. The child reparented // to init and is still training (or finished during the downtime), so we tail @@ -467,6 +464,13 @@ export async function runTraining({ jobId, runId, pythonPath = null, resumeCheck // staging, no validation, no fresh spawn. wireProcLifecycle then drives the // exact same line-handling/finalize path as a normal spawn, so a run that // completed mid-restart still registers its LoRA instead of being discarded. + // + // This runs BEFORE the fresh claimHeavyLocalJob() below: the survivor already + // holds the machine-wide accelerator claim from before the restart (handed + // off to its PID pre-crash), so acquiring a NEW claim here would see that + // live claim as a competing job and refuse it — failing every restart- + // survived run outright. Adopt the existing claim instead of contending for + // a fresh one. if (reattach) { const proc = await reattachDetached(join(dir, '.detached')); if (!proc) { @@ -475,12 +479,22 @@ export async function runTraining({ jobId, runId, pythonPath = null, resumeCheck // pre-#1332 reap path would have left it. return failBeforeSpawn('Trainer did not survive the restart — marking failed; resume from the latest checkpoint.'); } + heavyClaim = (await adoptHeavyLocalJob({ kind: 'LoRA training', id: jobId, pid: proc.pid })) + || (await claimHeavyLocalJob({ kind: 'LoRA training', id: jobId })); + if (!heavyClaim.ok) return failBeforeSpawn(heavyClaim.message); + // Only true on the claimHeavyLocalJob fallback (no matching on-disk claim + // survived) — adoptHeavyLocalJob only ever returns a claim already + // recorded against this exact pid. + if (heavyClaim.holder?.pid !== proc.pid) await heavyClaim.handoffTo?.(proc.pid); console.log(`🔁 training [${shortId(jobId)}] re-attached to surviving trainer pid ${proc.pid} (run ${shortId(runId)})`); trainingEvents.emit('status', { generationId: jobId, message: 'Re-attached to trainer that survived a restart' }); wireProcLifecycle(proc, { isReattach: true }); return; } + heavyClaim = await claimHeavyLocalJob({ kind: 'LoRA training', id: jobId }); + if (!heavyClaim.ok) return failBeforeSpawn(heavyClaim.message); + // Re-validate — the dataset may have been edited/deleted while queued. Skip // the caption identity-leak gate ONLY for a run that already opted past it: // one launched with an explicit "Train anyway" (persisted on the record) or a