Skip to content

Commit 0e3dcff

Browse files
fix(server): reclaim abandoned model downloads and retry GPU discovery
1 parent 26b9761 commit 0e3dcff

4 files changed

Lines changed: 55 additions & 7 deletions

File tree

‎apps/server/src/speech/SpeechModelStore.test.ts‎

Lines changed: 23 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -74,12 +74,24 @@ const harness = Effect.gen(function* () {
7474
);
7575
}),
7676
);
77-
const store = yield* SpeechModelStore.SpeechModelStore.pipe(
77+
const openStore = SpeechModelStore.SpeechModelStore.pipe(
7878
Effect.provide(
7979
SpeechModelStore.layer.pipe(Layer.provide(Layer.succeed(HttpClient.HttpClient, client))),
8080
),
8181
);
82-
return { model, directory, stale, fs, requests, bodies, transportFailure, bodyStarted, store };
82+
const store = yield* openStore;
83+
return {
84+
model,
85+
directory,
86+
stale,
87+
fs,
88+
requests,
89+
bodies,
90+
transportFailure,
91+
bodyStarted,
92+
store,
93+
openStore,
94+
};
8395
});
8496

8597
const dependencies = ServerConfig.layerTest("/tmp", { prefix: "speech-model-store-" }).pipe(
@@ -136,7 +148,7 @@ it.effect("rejects a same-size download whose SHA-256 differs from the catalog",
136148
}).pipe(Effect.provide(dependencies)),
137149
);
138150

139-
it.effect("cancelling a download removes the partial file and keeps the existing model", () =>
151+
it.effect("cleans only download partials on cancellation or restart", () =>
140152
Effect.gen(function* () {
141153
const h = yield* harness;
142154
h.bodies.current = () => Stream.concat(Stream.succeed(bytes.subarray(0, SPLIT)), Stream.never);
@@ -145,5 +157,13 @@ it.effect("cancelling a download removes the partial file and keeps the existing
145157
yield* Fiber.interrupt(fiber);
146158
expect(yield* h.fs.readDirectory(h.directory)).toEqual(["model.gguf"]);
147159
expect([...(yield* h.fs.readFile(h.stale))]).toEqual([...new Uint8Array(bytes.length)]);
160+
const orphan = `${h.stale}.f660de6c-e172-4146-8ffe-d3c32eb48cd4.part`;
161+
yield* h.fs.writeFile(orphan, bytes.subarray(0, SPLIT));
162+
yield* h.fs.writeFileString(`${h.stale}.notes.part`, "not a download");
163+
yield* h.openStore;
164+
expect((yield* h.fs.readDirectory(h.directory)).sort()).toEqual([
165+
"model.gguf",
166+
"model.gguf.notes.part",
167+
]);
148168
}).pipe(Effect.provide(dependencies)),
149169
);

‎apps/server/src/speech/SpeechModelStore.ts‎

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -54,6 +54,8 @@ export class SpeechModelStore extends Context.Service<
5454
>()("t3/speech/SpeechModelStore") {}
5555

5656
const DOWNLOAD_ATTEMPTS = 3;
57+
const PARTIAL_FILENAME =
58+
/\.gguf\.[0-9a-f]{8}-[0-9a-f]{4}-4[0-9a-f]{3}-[89ab][0-9a-f]{3}-[0-9a-f]{12}\.part$/i;
5759

5860
const make = Effect.gen(function* () {
5961
const fs = yield* FileSystem.FileSystem;
@@ -62,6 +64,12 @@ const make = Effect.gen(function* () {
6264
const random = yield* SecureRandom.SecureRandom;
6365
const config = yield* ServerConfig.ServerConfig;
6466
const directory = path.join(config.stateDir, "speech", "models");
67+
// The server already owns this state directory; no download can be active before the store opens.
68+
for (const name of yield* fs.readDirectory(directory).pipe(Effect.orElseSucceed(() => []))) {
69+
if (PARTIAL_FILENAME.test(name)) {
70+
yield* fs.remove(path.join(directory, name), { force: true }).pipe(Effect.ignore);
71+
}
72+
}
6573
const modelPath = (model: SpeechModel) => path.join(directory, model.filename);
6674
// Files whose digest failed, by identity, so status polls stop calling them ready.
6775
const corrupt = new Map<string, string>();

‎apps/server/src/speech/SpeechService.test.ts‎

Lines changed: 19 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -12,6 +12,7 @@ import * as Exit from "effect/Exit";
1212
import * as HostProcess from "@t3tools/shared/HostProcess";
1313
import { DEFAULT_SPEECH_TRANSCRIPTION_OPTIONS } from "@t3tools/contracts";
1414
import * as ServerSettings from "../serverSettings.ts";
15+
import { listNativeSpeechGpuDevices } from "./native.ts";
1516
import * as SpeechModelStore from "./SpeechModelStore.ts";
1617
import * as SpeechService from "./SpeechService.ts";
1718

@@ -398,6 +399,24 @@ it.effect("reports unsupported hosts without wrapping a synthetic cause", () =>
398399
}),
399400
);
400401

402+
it.effect.each([{ devices: [] }, { devices: [{ id: "gpu", name: "Test GPU" }] }])(
403+
"retries failed GPU discovery without repeatedly probing on status polls: $devices",
404+
({ devices }) =>
405+
Effect.gen(function* () {
406+
const speech = yield* SpeechService.SpeechService;
407+
const discover = vi.mocked(listNativeSpeechGpuDevices).mockReset();
408+
discover.mockRejectedValueOnce(new Error("driver not ready")).mockResolvedValueOnce(devices);
409+
expect(yield* speech.status).toMatchObject({ gpuDevices: [] });
410+
yield* speech.status;
411+
expect(discover).toHaveBeenCalledOnce();
412+
yield* TestClock.adjust("1 minute");
413+
expect(yield* speech.status).toMatchObject({ gpuDevices: devices });
414+
yield* TestClock.adjust("1 minute");
415+
expect(yield* speech.status).toMatchObject({ gpuDevices: devices });
416+
expect(discover).toHaveBeenCalledTimes(2);
417+
}).pipe(Effect.provide(layer)),
418+
);
419+
401420
it.effect(
402421
"retries Auto-placed inference with the dictionary and the CPU's vocabulary capability",
403422
() =>

‎apps/server/src/speech/SpeechService.ts‎

Lines changed: 5 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,7 @@ import * as HostProcess from "@t3tools/shared/HostProcess";
1010
import * as Cause from "effect/Cause";
1111
import * as Clock from "effect/Clock";
1212
import * as Context from "effect/Context";
13+
import * as Duration from "effect/Duration";
1314
import * as Effect from "effect/Effect";
1415
import * as Exit from "effect/Exit";
1516
import * as Fiber from "effect/Fiber";
@@ -183,7 +184,9 @@ const make = Effect.gen(function* () {
183184
let loadedAcceleration: string | undefined;
184185
let autoCpuModel: string | undefined;
185186
let loading: Promise<LoadedModel> | undefined;
186-
let gpuDevices: Promise<{ id: string; name: string }[]> | undefined;
187+
const discoverGpuDevices = yield* Effect.tryPromise(() => listNativeSpeechGpuDevices()).pipe(
188+
Effect.cachedWithTTL((exit) => (Exit.isSuccess(exit) ? Duration.infinity : "1 minute")),
189+
);
187190
let downloading:
188191
| {
189192
modelId: string;
@@ -438,9 +441,7 @@ const make = Effect.gen(function* () {
438441
supportsLanguageDetection: definition.supportsLanguageDetection,
439442
acceleration: settings.speechAcceleration,
440443
modelUnloadTimeout: settings.speechModelUnloadTimeout,
441-
gpuDevices: yield* Effect.promise(
442-
() => (gpuDevices ??= listNativeSpeechGpuDevices().catch(() => [])),
443-
),
444+
gpuDevices: yield* discoverGpuDevices.pipe(Effect.orElseSucceed(() => [])),
444445
};
445446
});
446447

0 commit comments

Comments
 (0)