954 lines · 32.2 KB · TypeScript
RawDownload
import { expect, test } from "bun:test";import { existsSync, mkdirSync, mkdtempSync, rmSync, writeFileSync } from "node:fs";import { tmpdir } from "node:os";import { join } from "node:path";import { Effect } from "effect";import { parseRecipe } from "../models/recipes/recipe-serializer";import type { GpuInfo, ProcessInfo, Recipe } from "../models/types";import { createGpuLeaseRegistry, type GpuLeaseRegistry } from "../system/gpu-leases";import type { ChatterboxRuntimeState } from "./runtime";import type { ChatterboxInstallOptions } from "./runtime";import {  SpeechService,  SpeechServiceError,  type SpeechDiskAvailability,  type SpeechEngineState,  type SpeechGpuLeaseGuard,  type SpeechRuntime,  type SpeechVoiceStore,  type SpeechWorker,} from "./service";import type { ChatterboxSynthesisInput, ChatterboxSynthesisResult } from "./worker-client";import type { VoiceProfile } from "./voice-store"; const PRO_UUIDS = [  "GPU-00000000-0000-0000-0000-000000000001",  "GPU-00000000-0000-0000-0000-000000000002",  "GPU-00000000-0000-0000-0000-000000000003",  "GPU-00000000-0000-0000-0000-000000000004",] as const;const SPEECH_UUID = "GPU-00000000-0000-0000-0000-000000003090";const INSTALLED: ChatterboxRuntimeState = {  status: "installed",  packageVersion: "0.1.7",  modelRevision: "749d1c1a46eb10492095d68fbcf55691ccf137cd",  gpuUuid: SPEECH_UUID,  installedAt: "2026-07-09T12:00:00.000Z",}; const gpu = (index: number, uuid: string, name: string): GpuInfo => ({  uuid,  index,  name,  memory_total_mb: 96_000,  memory_used_mb: 0,  memory_free_mb: 96_000,  utilization_pct: 0,  temp_c: 30,  power_draw: 0,  power_limit: 0,}); const gpus = (): GpuInfo[] => [  gpu(0, PRO_UUIDS[0], "NVIDIA RTX PRO 6000 Blackwell"),  gpu(1, PRO_UUIDS[1], "NVIDIA RTX PRO 6000 Blackwell"),  gpu(2, PRO_UUIDS[2], "NVIDIA RTX PRO 6000 Blackwell"),  gpu(3, SPEECH_UUID, "NVIDIA GeForce RTX 3090"),  gpu(4, PRO_UUIDS[3], "NVIDIA RTX PRO 6000 Blackwell"),]; const modelRecipe = (selector = "0,1,2,4"): Recipe =>  parseRecipe({    id: "speech-test-model",    name: "Speech test model",    model_path: "/models/speech-test",    env_vars: { CUDA_VISIBLE_DEVICES: selector },  }); const processInfo = (): ProcessInfo => ({  pid: 8123,  backend: "vllm",  model_path: "/models/speech-test",  port: 8000,  served_model_name: null,}); class FakeRuntime implements SpeechRuntime {  readonly paths: { readonly pythonPath: string };  installCalls = 0;  cancelCalls = 0;  cancelError: Error | null = null;  startError: Error | null = null;  private resolveInstall: ((state: ChatterboxRuntimeState) => void) | null = null;  private installPromise: Promise<ChatterboxRuntimeState> | null = null;   constructor(    directory: string,    private state: ChatterboxRuntimeState = INSTALLED,  ) {    this.paths = { pythonPath: join(directory, "runtime", "bin", "python") };  }   getState(): ChatterboxRuntimeState {    return this.state;  }   startInstall(gpuUuid: string, options: ChatterboxInstallOptions = {}): ChatterboxRuntimeState {    if (      this.state.status === "installing" ||      (this.state.status === "installed" && !options.repair)    ) {      return this.state;    }    this.installCalls += 1;    if (this.startError) throw this.startError;    this.state = {      status: "installing",      stage: "prefetching_model",      progress: 0.75,      gpuUuid,    };    this.installPromise = new Promise((resolve) => {      this.resolveInstall = resolve;    });    return this.state;  }   waitForInstall(): Promise<ChatterboxRuntimeState> {    return this.installPromise ?? Promise.resolve(this.state);  }   cancelInstall(): Promise<void> {    this.cancelCalls += 1;    if (this.state.status === "installing") {      this.state = {        status: "error",        gpuUuid: this.state.gpuUuid,        message: "Chatterbox install cancelled",      };      this.resolveInstall?.(this.state);      this.resolveInstall = null;    }    return this.cancelError ? Promise.reject(this.cancelError) : Promise.resolve();  }   finishInstall(): void {    this.state = INSTALLED;    this.resolveInstall?.(this.state);    this.resolveInstall = null;  }} class FakeVoiceStore implements SpeechVoiceStore {  readonly profiles: VoiceProfile[] = [    {      id: "voice_00000000000000000000000000000000",      name: "Sero",      duration_ms: 10_000,      created_at: "2026-07-09T12:00:00.000Z",    },  ];  plaintextError: Error | null = null;   constructor(private readonly referencePath: string) {}   list(): VoiceProfile[] {    return [...this.profiles];  }   create(input: {    name: string;    durationMs: number;    consent: string;    audio: Uint8Array;  }): Promise<VoiceProfile> {    const profile = {      id: "voice_11111111111111111111111111111111",      name: input.name,      duration_ms: input.durationMs,      created_at: "2026-07-09T13:00:00.000Z",    } satisfies VoiceProfile;    this.profiles.push(profile);    return Promise.resolve(profile);  }   delete(id: string): Promise<boolean> {    const index = this.profiles.findIndex((profile) => profile.id === id);    if (index < 0) return Promise.resolve(false);    this.profiles.splice(index, 1);    return Promise.resolve(true);  }   withPlaintext<A>(_id: string, use: (path: string) => Promise<A>): Promise<A> {    if (this.plaintextError) return Promise.reject(this.plaintextError);    return use(this.referencePath);  }} const wave = (): Buffer => {  const bytes = Buffer.alloc(46);  bytes.write("RIFF", 0, "ascii");  bytes.writeUInt32LE(38, 4);  bytes.write("WAVE", 8, "ascii");  bytes.write("fmt ", 12, "ascii");  bytes.writeUInt32LE(16, 16);  bytes.writeUInt16LE(1, 20);  bytes.writeUInt16LE(1, 22);  bytes.writeUInt32LE(24_000, 24);  bytes.writeUInt32LE(48_000, 28);  bytes.writeUInt16LE(2, 32);  bytes.writeUInt16LE(16, 34);  bytes.write("data", 36, "ascii");  bytes.writeUInt32LE(2, 40);  return bytes;}; class FakeWorker implements SpeechWorker {  shutdownError: Error | null = null;  settleError: Error | null = null;  terminateError: Error | null = null;  terminateCalls = 0;  synthesizeStarted = 0;  private blocker: Promise<void> | null = null;  private releaseBlocker: (() => void) | null = null;  private rejectBlocker: ((error: Error) => void) | null = null;  outputBytes: Uint8Array = wave();   constructor(private readonly outputPath: string) {}   blockSynthesis(): void {    this.blocker = new Promise((resolve, reject) => {      this.releaseBlocker = resolve;      this.rejectBlocker = reject;    });  }   releaseSynthesis(): void {    this.releaseBlocker?.();    this.releaseBlocker = null;    this.rejectBlocker = null;    this.blocker = null;  }   async synthesize(_input: ChatterboxSynthesisInput): Promise<ChatterboxSynthesisResult> {    this.synthesizeStarted += 1;    await this.blocker;    mkdirSync(join(this.outputPath, ".."), { recursive: true });    writeFileSync(this.outputPath, this.outputBytes);    return { path: this.outputPath, sampleRate: 24_000 };  }   shutdown(): Promise<void> {    return this.shutdownError ? Promise.reject(this.shutdownError) : Promise.resolve();  }   settleTermination(): Promise<void> {    return this.settleError ? Promise.reject(this.settleError) : Promise.resolve();  }   terminate(): Promise<void> {    this.terminateCalls += 1;    if (this.terminateError) return Promise.reject(this.terminateError);    this.rejectBlocker?.(new Error("Speech operation was interrupted"));    this.rejectBlocker = null;    this.releaseBlocker = null;    this.blocker = null;    return Promise.resolve();  }} const stableEngine = (recipe: Recipe | null = modelRecipe()): SpeechEngineState => ({  getCurrentProcess: async () => processInfo(),  getCurrentRecipe: async () => recipe,}); const noModelEngine = (): SpeechEngineState => ({  getCurrentProcess: async () => null,  getCurrentRecipe: async () => null,}); const readyDisk = (): SpeechDiskAvailability => ({  totalBytes: 256 * 1024 ** 3,  availableBytes: 80 * 1024 ** 3,}); const fixture = (  overrides: {    runtime?: SpeechRuntime;    engine?: SpeechEngineState;    registry?: GpuLeaseRegistry;    workerFactory?: (lease: SpeechGpuLeaseGuard) => SpeechWorker;    diskAvailability?: () => SpeechDiskAvailability | null;    gpuInfo?: () => GpuInfo[];    environment?: NodeJS.ProcessEnv;    normalizeReference?: (      input: Uint8Array,      dataDirectory: string,    ) => Promise<{ audio: Uint8Array; durationMs: number }>;    computeGpuUuids?: () => Promise<readonly string[]>;  } = {},): {  directory: string;  registry: GpuLeaseRegistry;  runtime: SpeechRuntime;  worker: FakeWorker;  voiceStore: FakeVoiceStore;  service: SpeechService;} => {  const directory = mkdtempSync(join(tmpdir(), "local-studio-speech-service-"));  const output = join(directory, "runtime", "speech", "outputs", "result.wav");  const reference = join(directory, "reference.wav");  writeFileSync(reference, wave());  const registry = overrides.registry ?? createGpuLeaseRegistry();  const runtime = overrides.runtime ?? new FakeRuntime(directory);  const worker = new FakeWorker(output);  const voiceStore = new FakeVoiceStore(reference);  const service = new SpeechService({    dataDirectory: directory,    databasePath: join(directory, "controller.db"),    engine: overrides.engine ?? stableEngine(),    gpuLeaseRegistry: registry,    gpuInfo: overrides.gpuInfo ?? gpus,    runtime,    voiceStore,    workerFactory: overrides.workerFactory ?? ((): SpeechWorker => worker),    diskAvailability: overrides.diskAvailability ?? readyDisk,    environment: overrides.environment ?? {},    ...(overrides.normalizeReference ? { normalizeReference: overrides.normalizeReference } : {}),    computeGpuUuids: overrides.computeGpuUuids ?? (async (): Promise<readonly string[]> => []),    resolveBinary: (name): string | null => (name === "uv" ? "/opt/uv" : null),  });  return { directory, registry, runtime, worker, voiceStore, service };}; const speechFailure = async (promise: Promise<unknown>): Promise<SpeechServiceError> => {  try {    await promise;  } catch (error) {    if (error instanceof SpeechServiceError) return error;    throw error;  }  throw new Error("Expected speech operation to fail");}; const waitFor = async (predicate: () => boolean): Promise<void> => {  for (let attempt = 0; attempt < 100; attempt += 1) {    if (predicate()) return;    await Promise.resolve();  }  throw new Error("Expected test condition was not reached");}; test("seeds model leases and presents a live speech lease before spawning the worker", async () => {  const factoryLeases: SpeechGpuLeaseGuard[] = [];  const created = fixture({    workerFactory: (lease): SpeechWorker => {      factoryLeases.push(lease);      return created.worker;    },  });  try {    const output = await created.service.synthesize({      text: "Local speech is isolated",      voiceId: "voice_00000000000000000000000000000000",    });     expect(output.audio).toEqual(wave());    expect(factoryLeases[0]?.uuid).toBe(SPEECH_UUID);    expect(await Effect.runPromise(created.registry.snapshot())).toEqual([      ...PRO_UUIDS.map((uuid) => ({ uuid, owner: "llm" as const })),      { uuid: SPEECH_UUID, owner: "speech" },    ]);    expect(existsSync(join(created.directory, "runtime", "speech", "outputs", "result.wav"))).toBe(      false,    );  } finally {    await created.service.stop().catch(() => undefined);    rmSync(created.directory, { recursive: true, force: true });  }}); test("blocks unknown and unresolved model processes before claiming speech", async () => {  const unknown = fixture({ engine: stableEngine(null) });  const unresolved = fixture({ engine: stableEngine(modelRecipe("0,missing")) });  try {    expect(      (        await speechFailure(          unknown.service.synthesize({            text: "Blocked",            voiceId: "voice_00000000000000000000000000000000",          }),        )      ).code,    ).toBe("model_process_unknown");    expect(      (        await speechFailure(          unresolved.service.synthesize({            text: "Blocked",            voiceId: "voice_00000000000000000000000000000000",          }),        )      ).code,    ).toBe("model_gpu_unresolved");    expect(await Effect.runPromise(unknown.registry.snapshot())).toEqual([]);    expect(await Effect.runPromise(unresolved.registry.snapshot())).toEqual([]);  } finally {    rmSync(unknown.directory, { recursive: true, force: true });    rmSync(unresolved.directory, { recursive: true, force: true });  }}); test("accepts an explicit full UUID and fails closed when UUID telemetry is missing", async () => {  const renamedGpus = (): GpuInfo[] =>    gpus().map((candidate) =>      candidate.uuid === SPEECH_UUID        ? { ...candidate, name: "NVIDIA dedicated speech GPU" }        : candidate,    );  const configured = fixture({    engine: noModelEngine(),    gpuInfo: renamedGpus,    environment: { LOCAL_STUDIO_SPEECH_GPU_UUID: SPEECH_UUID },  });  const missing = fixture({    engine: noModelEngine(),    gpuInfo: () =>      gpus().map((candidate) => {        const withoutUuid = { ...candidate };        delete withoutUuid.uuid;        return withoutUuid;      }),  });  try {    expect(configured.service.getStatus().gpu?.uuid).toBe(SPEECH_UUID);    await configured.service.synthesize({      text: "Explicit target",      voiceId: "voice_00000000000000000000000000000000",    });    expect(missing.service.getStatus().gpu).toBeNull();    expect(      (        await speechFailure(          missing.service.synthesize({            text: "No telemetry",            voiceId: "voice_00000000000000000000000000000000",          }),        )      ).code,    ).toBe("speech_gpu_telemetry_missing");    expect(await Effect.runPromise(missing.registry.snapshot())).toEqual([]);  } finally {    await configured.service.stop().catch(() => undefined);    rmSync(configured.directory, { recursive: true, force: true });    rmSync(missing.directory, { recursive: true, force: true });  }}); test("does not steal an in-flight model lease before its process appears", async () => {  const registry = createGpuLeaseRegistry();  await Effect.runPromise(registry.claim("llm", [SPEECH_UUID]));  const created = fixture({ registry, engine: noModelEngine() });  try {    const error = await speechFailure(      created.service.synthesize({        text: "Blocked",        voiceId: "voice_00000000000000000000000000000000",      }),    );    expect(error.code).toBe("model_gpu_transition");    expect(await Effect.runPromise(registry.snapshot())).toEqual([      { uuid: SPEECH_UUID, owner: "llm" },    ]);  } finally {    rmSync(created.directory, { recursive: true, force: true });  }}); test("blocks an orphan compute process before reserving the speech GPU", async () => {  let workerCreations = 0;  const created = fixture({    engine: noModelEngine(),    computeGpuUuids: async () => [SPEECH_UUID],    workerFactory: (): SpeechWorker => {      workerCreations += 1;      return created.worker;    },  });  try {    const error = await speechFailure(      created.service.synthesize({        text: "Do not overlap this process",        voiceId: "voice_00000000000000000000000000000000",      }),    );    expect(error.code).toBe("speech_gpu_compute_busy");    expect(workerCreations).toBe(0);    expect(await Effect.runPromise(created.registry.snapshot())).toEqual([]);  } finally {    rmSync(created.directory, { recursive: true, force: true });  }}); test("releases admission when the post-lease compute query fails", async () => {  let queries = 0;  const created = fixture({    engine: noModelEngine(),    computeGpuUuids: async () => {      queries += 1;      if (queries === 1) return [];

Lines 1–500 of 954.