diff --git a/packages/server/src/server/speech/providers/local/models.test.ts b/packages/server/src/server/speech/providers/local/models.test.ts new file mode 100644 index 0000000000..e40cc25384 --- /dev/null +++ b/packages/server/src/server/speech/providers/local/models.test.ts @@ -0,0 +1,70 @@ +import { execFileSync } from "node:child_process"; +import { + existsSync, + mkdirSync, + mkdtempSync, + readFileSync, + rmSync, + truncateSync, + writeFileSync, +} from "node:fs"; +import { tmpdir } from "node:os"; +import path from "node:path"; +import pino from "pino"; +import { describe, expect, onTestFinished, test } from "vitest"; + +import { + ensureLocalSpeechModels, + getLocalSpeechModelDir, + listLocalSpeechModels, + listMissingLocalSpeechModels, +} from "./models.js"; + +const logger = pino({ level: "silent" }); +const modelId = "parakeet-tdt-0.6b-v2-int8"; +const spec = listLocalSpeechModels().find((model) => model.id === modelId)!; + +function writeCompleteModel(modelDir: string): void { + mkdirSync(modelDir, { recursive: true }); + for (const file of spec.requiredFiles) { + writeFileSync(path.join(modelDir, file), `complete ${file}`.padEnd(64_000, ".")); + } +} + +function setupModelsDirWithArchive(): { modelsDir: string; modelDir: string; archive: string } { + const root = mkdtempSync(path.join(tmpdir(), "paseo-local-models-")); + onTestFinished(() => rmSync(root, { recursive: true, force: true })); + const sourceRoot = path.join(root, "source"); + writeCompleteModel(path.join(sourceRoot, spec.extractedDir)); + const modelsDir = path.join(root, "models"); + const archive = path.join( + modelsDir, + ".downloads", + path.basename(new URL(spec.archiveUrl).pathname), + ); + mkdirSync(path.dirname(archive), { recursive: true }); + execFileSync("tar", ["cf", archive, "-C", sourceRoot, spec.extractedDir]); + return { modelsDir, modelDir: getLocalSpeechModelDir(modelsDir, modelId), archive }; +} + +describe("local speech models", () => { + test("reinstalls a model whose extraction was interrupted", async () => { + // An interrupted extraction leaves the archive in .downloads and a truncated file in place. + const { modelsDir, modelDir, archive } = setupModelsDirWithArchive(); + writeCompleteModel(modelDir); + const encoder = path.join(modelDir, spec.requiredFiles[0]); + truncateSync(encoder, 1_000); + + expect(await listMissingLocalSpeechModels({ modelsDir, modelIds: [modelId] })).toEqual([ + modelId, + ]); + + await ensureLocalSpeechModels({ modelsDir, modelIds: [modelId], logger }); + + expect(readFileSync(encoder, "utf8")).toBe( + `complete ${spec.requiredFiles[0]}`.padEnd(64_000, "."), + ); + expect(existsSync(archive)).toBe(false); + expect(await listMissingLocalSpeechModels({ modelsDir, modelIds: [modelId] })).toEqual([]); + }); +}); diff --git a/packages/server/src/server/speech/providers/local/models.ts b/packages/server/src/server/speech/providers/local/models.ts index f6c7cc4e4e..a884be86d7 100644 --- a/packages/server/src/server/speech/providers/local/models.ts +++ b/packages/server/src/server/speech/providers/local/models.ts @@ -1,4 +1,8 @@ -import { ensureSherpaOnnxModels, getSherpaOnnxModelDir } from "./sherpa/model-downloader.js"; +import { + ensureSherpaOnnxModels, + getSherpaOnnxModelDir, + isSherpaOnnxModelInstalled, +} from "./sherpa/model-downloader.js"; import { DEFAULT_LOCAL_STT_MODEL, DEFAULT_LOCAL_TTS_MODEL, @@ -30,6 +34,16 @@ export function getLocalSpeechModelDir(modelsDir: string, modelId: LocalSpeechMo return getSherpaOnnxModelDir(modelsDir, modelId); } +export async function listMissingLocalSpeechModels(options: { + modelsDir: string; + modelIds: LocalSpeechModelId[]; +}): Promise { + const installed = await Promise.all( + options.modelIds.map((modelId) => isSherpaOnnxModelInstalled(options.modelsDir, modelId)), + ); + return options.modelIds.filter((_modelId, index) => !installed[index]); +} + export async function ensureLocalSpeechModels(options: { modelsDir: string; modelIds: LocalSpeechModelId[]; diff --git a/packages/server/src/server/speech/providers/local/sherpa/model-downloader.test.ts b/packages/server/src/server/speech/providers/local/sherpa/model-downloader.test.ts index 993c959b4d..bc2e462a53 100644 --- a/packages/server/src/server/speech/providers/local/sherpa/model-downloader.test.ts +++ b/packages/server/src/server/speech/providers/local/sherpa/model-downloader.test.ts @@ -1,10 +1,21 @@ import { describe, expect, test } from "vitest"; -import { mkdtempSync, mkdirSync, writeFileSync } from "node:fs"; +import { + mkdtempSync, + mkdirSync, + writeFileSync, + readFileSync, + rmSync, + existsSync, + truncateSync, + readdirSync, +} from "node:fs"; +import { execFileSync } from "node:child_process"; import { tmpdir } from "node:os"; import path from "node:path"; import pino from "pino"; import { ensureSherpaOnnxModel, getSherpaOnnxModelDir } from "./model-downloader.js"; +import { getSherpaOnnxModelSpec } from "./model-catalog.js"; function makeTmpDir(): string { return mkdtempSync(path.join(tmpdir(), "paseo-speech-models-")); @@ -13,6 +24,67 @@ function makeTmpDir(): string { const logger = pino({ level: "silent" }); describe("sherpa model downloader", () => { + test("does not expose a partial model when tar extraction fails", async () => { + const root = makeTmpDir(); + const modelsDir = path.join(root, "models"); + const spec = getSherpaOnnxModelSpec("kokoro-en-v0_19"); + const sourceDir = path.join(root, "source", spec.extractedDir); + const archive = path.join( + modelsDir, + ".downloads", + path.basename(new URL(spec.archiveUrl).pathname), + ); + try { + mkdirSync(sourceDir, { recursive: true }); + mkdirSync(path.dirname(archive), { recursive: true }); + writeFileSync(path.join(sourceDir, "model.onnx"), Buffer.alloc(32_768, 1)); + execFileSync("tar", [ + "cf", + archive, + "-C", + path.dirname(sourceDir), + `${spec.extractedDir}/model.onnx`, + ]); + truncateSync(archive, 1_024); + await expect(ensureSherpaOnnxModel({ modelsDir, modelId: spec.id, logger })).rejects.toThrow( + "tar exited", + ); + expect(existsSync(getSherpaOnnxModelDir(modelsDir, spec.id))).toBe(false); + expect(existsSync(archive)).toBe(true); + expect(readdirSync(modelsDir)).toEqual([".downloads"]); + } finally { + rmSync(root, { recursive: true, force: true }); + } + }); + + test("repairs a partial legacy extraction when its cached archive remains", async () => { + const root = makeTmpDir(); + const modelsDir = path.join(root, "models"); + const spec = getSherpaOnnxModelSpec("kokoro-en-v0_19"); + const sourceDir = path.join(root, "source", spec.extractedDir); + const modelDir = getSherpaOnnxModelDir(modelsDir, spec.id); + const archive = path.join( + modelsDir, + ".downloads", + path.basename(new URL(spec.archiveUrl).pathname), + ); + try { + for (const dir of [sourceDir, modelDir]) { + mkdirSync(path.join(dir, "espeak-ng-data"), { recursive: true }); + for (const filename of ["model.onnx", "voices.bin", "tokens.txt"]) { + writeFileSync(path.join(dir, filename), "complete model data"); + } + } + mkdirSync(path.dirname(archive), { recursive: true }); + execFileSync("tar", ["cf", archive, "-C", path.dirname(sourceDir), spec.extractedDir]); + writeFileSync(path.join(modelDir, "model.onnx"), "truncated"); + await ensureSherpaOnnxModel({ modelsDir, modelId: spec.id, logger }); + expect(readFileSync(path.join(modelDir, "model.onnx"), "utf8")).toBe("complete model data"); + } finally { + rmSync(root, { recursive: true, force: true }); + } + }); + test("getSherpaOnnxModelDir maps modelId to extractedDir", () => { const modelsDir = "/tmp/models"; expect(getSherpaOnnxModelDir(modelsDir, "parakeet-tdt-0.6b-v2-int8")).toContain( diff --git a/packages/server/src/server/speech/providers/local/sherpa/model-downloader.ts b/packages/server/src/server/speech/providers/local/sherpa/model-downloader.ts index 8ce21f5e01..cdd286389f 100644 --- a/packages/server/src/server/speech/providers/local/sherpa/model-downloader.ts +++ b/packages/server/src/server/speech/providers/local/sherpa/model-downloader.ts @@ -90,6 +90,38 @@ async function extractTarArchive( }); } +interface SherpaOnnxModelPaths { + modelDir: string; + archivePath: string; + extractionDir: string; +} + +function getSherpaOnnxModelPaths( + modelsDir: string, + modelId: SherpaOnnxModelId, +): SherpaOnnxModelPaths { + const spec = getSherpaOnnxModelSpec(modelId); + const downloadsDir = path.join(modelsDir, ".downloads"); + return { + modelDir: path.join(modelsDir, spec.extractedDir), + archivePath: path.join(downloadsDir, path.basename(new URL(spec.archiveUrl).pathname)), + extractionDir: path.join(downloadsDir, `${spec.extractedDir}.extracting`), + }; +} + +// The archive is removed only as its model moves into place, so a model whose archive is +// still in .downloads did not finish installing and may hold truncated files. +export async function isSherpaOnnxModelInstalled( + modelsDir: string, + modelId: SherpaOnnxModelId, +): Promise { + const spec = getSherpaOnnxModelSpec(modelId); + const { modelDir, archivePath } = getSherpaOnnxModelPaths(modelsDir, modelId); + return ( + !(await isNonEmptyFile(archivePath)) && (await hasRequiredFiles(modelDir, spec.requiredFiles)) + ); +} + async function isNonEmptyFile(filePath: string): Promise { try { const s = await stat(filePath); @@ -110,18 +142,17 @@ export async function ensureSherpaOnnxModel( }); const spec = getSherpaOnnxModelSpec(options.modelId); - const modelDir = path.join(options.modelsDir, spec.extractedDir); - if (await hasRequiredFiles(modelDir, spec.requiredFiles)) { + const { modelDir, archivePath, extractionDir } = getSherpaOnnxModelPaths( + options.modelsDir, + options.modelId, + ); + if (await isSherpaOnnxModelInstalled(options.modelsDir, options.modelId)) { return modelDir; } logger.info({ modelsDir: options.modelsDir }, "Starting model download"); try { - const downloadsDir = path.join(options.modelsDir, ".downloads"); - const archiveFilename = path.basename(new URL(spec.archiveUrl).pathname); - const archivePath = path.join(downloadsDir, archiveFilename); - if (!(await isNonEmptyFile(archivePath))) { await downloadToFile({ url: spec.archiveUrl, @@ -138,32 +169,40 @@ export async function ensureSherpaOnnxModel( }, "Extracting model archive", ); - await extractTarArchive(archivePath, options.modelsDir, options.signal); - - logger.info( - { - modelId: options.modelId, - modelDir, - }, - "Verifying downloaded model files", - ); - if (!(await hasRequiredFiles(modelDir, spec.requiredFiles))) { - throw new Error( - `Downloaded and extracted ${archiveFilename}, but required files are still missing in ${modelDir}.`, - ); - } - - logger.info( - { - modelId: options.modelId, - archivePath, - }, - "Finalizing model artifacts", - ); + // Extract beside the archive and move into place once complete, so an interrupted + // extraction never leaves partial files at modelDir. + await rm(extractionDir, { recursive: true, force: true }); try { - await rm(archivePath, { force: true }); - } catch { - // ignore + await extractTarArchive(archivePath, extractionDir, options.signal); + const extractedModelDir = path.join(extractionDir, spec.extractedDir); + + logger.info( + { + modelId: options.modelId, + modelDir, + }, + "Verifying downloaded model files", + ); + if (!(await hasRequiredFiles(extractedModelDir, spec.requiredFiles))) { + throw new Error( + `Downloaded and extracted ${path.basename(archivePath)}, but required files are missing.`, + ); + } + + logger.info( + { + modelId: options.modelId, + archivePath, + }, + "Finalizing model artifacts", + ); + // A retained archive marks the model as not installed, so remove it before the model + // is moved into place. If removal fails, nothing at modelDir looks installed. + await rm(modelDir, { recursive: true, force: true }); + await rm(archivePath, { force: true, maxRetries: 3 }); + await rename(extractedModelDir, modelDir); + } finally { + await rm(extractionDir, { recursive: true, force: true }); } logger.info({ modelDir }, "Model download completed"); diff --git a/packages/server/src/server/speech/speech-runtime.test.ts b/packages/server/src/server/speech/speech-runtime.test.ts index 428ecd27aa..c9c21b3e89 100644 --- a/packages/server/src/server/speech/speech-runtime.test.ts +++ b/packages/server/src/server/speech/speech-runtime.test.ts @@ -36,8 +36,7 @@ vi.mock("./providers/openai/runtime.js", () => ({ vi.mock("./providers/local/models.js", () => ({ ensureLocalSpeechModels: ensureLocalSpeechModelsMock, - getLocalSpeechModelDir: vi.fn(() => ""), - listLocalSpeechModels: vi.fn(() => []), + listMissingLocalSpeechModels: vi.fn(async ({ modelIds }: { modelIds: string[] }) => modelIds), })); function createStubStt(id: string): SpeechToTextProvider { diff --git a/packages/server/src/server/speech/speech-runtime.ts b/packages/server/src/server/speech/speech-runtime.ts index 6e56a9e955..932f87c077 100644 --- a/packages/server/src/server/speech/speech-runtime.ts +++ b/packages/server/src/server/speech/speech-runtime.ts @@ -1,14 +1,8 @@ -import { stat } from "node:fs/promises"; -import { join } from "node:path"; import type { Logger } from "pino"; import type { PaseoOpenAIConfig, PaseoSpeechConfig } from "../bootstrap.js"; import type { LocalSpeechModelId } from "./providers/local/config.js"; -import { - ensureLocalSpeechModels, - getLocalSpeechModelDir, - listLocalSpeechModels, -} from "./providers/local/models.js"; +import { ensureLocalSpeechModels, listMissingLocalSpeechModels } from "./providers/local/models.js"; import { initializeLocalSpeechServices } from "./providers/local/runtime.js"; import { getOpenAiSpeechAvailability, @@ -76,48 +70,6 @@ function resolveRequestedSpeechProviders( }; } -async function hasRequiredLocalModelFile(filePath: string): Promise { - try { - const fileStat = await stat(filePath); - if (fileStat.isDirectory()) { - return true; - } - return fileStat.isFile() && fileStat.size > 0; - } catch { - return false; - } -} - -async function findMissingRequiredLocalModels(params: { - modelsDir: string | null; - requiredModelIds: LocalSpeechModelId[]; -}): Promise { - const { modelsDir, requiredModelIds } = params; - if (!modelsDir || requiredModelIds.length === 0) { - return []; - } - - const specsById = new Map(listLocalSpeechModels().map((model) => [model.id, model])); - const missing = new Set(); - - const checks = await Promise.all( - requiredModelIds.map(async (modelId) => { - const spec = specsById.get(modelId); - if (!spec) return { modelId, missing: true }; - const modelDir = getLocalSpeechModelDir(modelsDir, modelId); - const filePresence = await Promise.all( - spec.requiredFiles.map((relPath) => hasRequiredLocalModelFile(join(modelDir, relPath))), - ); - return { modelId, missing: !filePresence.every((present) => present) }; - }), - ); - for (const check of checks) { - if (check.missing) missing.add(check.modelId); - } - - return Array.from(missing); -} - function joinModelIds(modelIds: LocalSpeechModelId[]): string { if (modelIds.length === 0) { return "none"; @@ -493,10 +445,12 @@ export function createSpeechService(params: { }; const refreshMissingLocalModels = async (): Promise => { - missingLocalModelIds = await findMissingRequiredLocalModels({ - modelsDir: localModelConfig?.modelsDir ?? null, - requiredModelIds: localModelConfig?.defaultModelIds ?? [], - }); + missingLocalModelIds = localModelConfig + ? await listMissingLocalSpeechModels({ + modelsDir: localModelConfig.modelsDir, + modelIds: localModelConfig.defaultModelIds, + }) + : []; }; const reconcileServices = async (): Promise => {