import { expect, test } from "bun:test";import { mkdir, mkdtemp, rm, utimes, writeFile } from "node:fs/promises";import { tmpdir } from "node:os";import { join } from "node:path";import { Effect } from "effect";import { asRecipeId, type GpuInfo, type Recipe } from "../models/types";import { createGpuLeaseRegistry, GpuLeaseConflict, GpuLeaseLockFailure, resolveRecipeGpuUuids, type GpuLease,} from "./gpu-leases"; const proUuids = [ "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 rtxUuid = "GPU-00000000-0000-0000-0000-000000003090"; const temporaryLockDirectory = (): Promise<string> => mkdtemp(join(tmpdir(), "local-studio-gpu-leases-test-")); const lockPath = (directory: string, uuid: string): string => join(directory, `${uuid.toLowerCase()}.lock`); const writeHostLock = async ( directory: string, uuid: string, pid: number, processStartToken: string | null,): Promise<void> => { await mkdir(directory, { recursive: true, mode: 0o700 }); await writeFile( lockPath(directory, uuid), JSON.stringify({ version: 1, uuid, owner: "speech", pid, processStartToken, registryId: "stale-registry", }), { mode: 0o600 }, );}; function gpu(index: number, name: string, uuid: string): GpuInfo { return { 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, };} function recipe( environmentVariables: Record<string, string> | null = null, extraArguments: Record<string, unknown> = {},): Recipe { return { id: asRecipeId("gpu-lease-test"), name: "GPU lease test", model_path: "model", vision: null, backend: "vllm", runtime: { kind: "managed_venv", ref: "vllm-latest" }, env_vars: environmentVariables, tensor_parallel_size: 1, pipeline_parallel_size: 1, max_model_len: 4096, gpu_memory_utilization: 0.9, kv_cache_dtype: "auto", max_num_seqs: 1, trust_remote_code: false, tool_call_parser: null, reasoning_parser: null, enable_auto_tool_choice: false, quantization: null, dtype: null, host: "127.0.0.1", port: 8000, served_model_name: null, python_path: null, extra_args: extraArguments, max_thinking_tokens: null, thinking_mode: "auto", };} function fiveGpuHost(): GpuInfo[] { return [ ...proUuids.slice(0, 3).map((uuid, index) => gpu(index, "NVIDIA RTX PRO 6000 Blackwell", uuid)), gpu(3, "NVIDIA GeForce RTX 3090", rtxUuid), gpu(4, "NVIDIA RTX PRO 6000 Blackwell", proUuids[3]), ];} test("resolves the four PRO recipe without leasing the 3090", () => { const resolution = resolveRecipeGpuUuids( recipe({ CUDA_VISIBLE_DEVICES: "0,1,2,4" }), fiveGpuHost(), ); expect(resolution).toEqual({ source: "recipe", selector: "0,1,2,4", uuids: [...proUuids], unresolvedTokens: [], });}); test("reports an atomic all-GPU lease conflict", async () => { const registry = createGpuLeaseRegistry(); const allUuids = resolveRecipeGpuUuids(recipe(), fiveGpuHost()).uuids; await Effect.runPromise(registry.claim("llm", proUuids)); await Effect.runPromise(registry.claim("speech", [rtxUuid])); try { await Effect.runPromise(registry.replace("llm", allUuids)); throw new Error("expected a GPU lease conflict"); } catch (error) { expect(error).toBeInstanceOf(GpuLeaseConflict); if (error instanceof GpuLeaseConflict) { expect(error.requestedBy).toBe("llm"); expect(error.conflicts).toEqual([{ uuid: rtxUuid, heldBy: "speech" }]); } } const unchanged: GpuLease[] = [ ...proUuids.map((uuid): GpuLease => ({ uuid, owner: "llm" })), { uuid: rtxUuid, owner: "speech" }, ]; expect(await Effect.runPromise(registry.snapshot())).toEqual(unchanged);}); test("keeps UUID selectors stable when GPU indices reorder", () => { const reordered = [ gpu(0, "NVIDIA RTX PRO 6000 Blackwell", proUuids[3]), gpu(4, "NVIDIA RTX PRO 6000 Blackwell", proUuids[0]), ]; const selector = `${proUuids[0]},${proUuids[3]}`; expect( resolveRecipeGpuUuids(recipe({ CUDA_VISIBLE_DEVICES: selector }), reordered).uuids, ).toEqual([proUuids[0], proUuids[3]]);}); test("reports unresolved aliased selectors without falling back", () => { const resolution = resolveRecipeGpuUuids( recipe(null, { "cuda-visible-devices": `${proUuids[1]},99` }), fiveGpuHost(), ); expect(resolution.uuids).toEqual([proUuids[1]]); expect(resolution.unresolvedTokens).toEqual(["99"]);}); test("replaces and releases only the requesting owner leases", async () => { const registry = createGpuLeaseRegistry(); await Effect.runPromise(registry.claim("llm", [proUuids[0], proUuids[1]])); await Effect.runPromise(registry.claim("speech", [rtxUuid])); await Effect.runPromise(registry.replace("llm", [proUuids[2], proUuids[3]])); await Effect.runPromise(registry.release("speech")); await Effect.runPromise(registry.release("llm", [proUuids[2]])); expect(await Effect.runPromise(registry.snapshot())).toEqual([ { uuid: proUuids[3], owner: "llm" }, ]);}); test("conflicts across registries in the same live process", async () => { const directory = await temporaryLockDirectory(); const first = createGpuLeaseRegistry({ lockDirectory: directory }); const second = createGpuLeaseRegistry({ lockDirectory: directory }); try { await Effect.runPromise(first.claim("speech", [rtxUuid])); expect(await Effect.runPromise(first.claim("speech", [rtxUuid]))).toEqual([ { uuid: rtxUuid, owner: "speech" }, ]); try { await Effect.runPromise(second.claim("speech", [rtxUuid])); throw new Error("expected a cross-registry GPU lease conflict"); } catch (error) { expect(error).toBeInstanceOf(GpuLeaseConflict); if (error instanceof GpuLeaseConflict) { expect(error.conflicts).toEqual([{ uuid: rtxUuid, heldBy: "speech" }]); } } await Effect.runPromise(first.release("speech")); expect(await Effect.runPromise(second.claim("speech", [rtxUuid]))).toEqual([ { uuid: rtxUuid, owner: "speech" }, ]); } finally { await Effect.runPromise(first.release("speech")); await Effect.runPromise(second.release("speech")); await rm(directory, { recursive: true, force: true }); }}); test("reclaims a lock from a dead process", async () => { const directory = await temporaryLockDirectory(); const registry = createGpuLeaseRegistry({ lockDirectory: directory }); try { await writeHostLock(directory, rtxUuid, 2_147_483_647, "1"); expect(await Effect.runPromise(registry.claim("speech", [rtxUuid]))).toEqual([ { uuid: rtxUuid, owner: "speech" }, ]); } finally { await Effect.runPromise(registry.release("speech")); await rm(directory, { recursive: true, force: true }); }}); test("reclaims a reused Linux PID by process start token", async () => { if (process.platform !== "linux") return; const directory = await temporaryLockDirectory(); const registry = createGpuLeaseRegistry({ lockDirectory: directory }); try { await writeHostLock(directory, rtxUuid, process.pid, "0"); expect(await Effect.runPromise(registry.claim("speech", [rtxUuid]))).toEqual([ { uuid: rtxUuid, owner: "speech" }, ]); } finally { await Effect.runPromise(registry.release("speech")); await rm(directory, { recursive: true, force: true }); }}); test("fails closed on an unverifiable host lock record", async () => { const directory = await temporaryLockDirectory(); const registry = createGpuLeaseRegistry({ lockDirectory: directory }); try { await writeFile(lockPath(directory, rtxUuid), "not-json", { mode: 0o600 }); await expect(Effect.runPromise(registry.claim("speech", [rtxUuid]))).rejects.toBeInstanceOf( GpuLeaseLockFailure, ); } finally { await rm(directory, { recursive: true, force: true }); }}); test("recovers a stale reaper left during dead-lock reclamation", async () => { const directory = await temporaryLockDirectory(); const registry = createGpuLeaseRegistry({ lockDirectory: directory }); const reaper = `${lockPath(directory, rtxUuid)}.reaper`; try { await writeHostLock(directory, rtxUuid, 2_147_483_647, "1"); await mkdir(reaper, { mode: 0o700 }); const old = new Date(Date.now() - 10_000); await utimes(reaper, old, old); expect(await Effect.runPromise(registry.claim("speech", [rtxUuid]))).toEqual([ { uuid: rtxUuid, owner: "speech" }, ]); } finally { await Effect.runPromise(registry.release("speech")); await rm(directory, { recursive: true, force: true }); }}); test("rolls back new host locks before a failed replacement", async () => { const directory = await temporaryLockDirectory(); const primary = createGpuLeaseRegistry({ lockDirectory: directory }); const blocker = createGpuLeaseRegistry({ lockDirectory: directory }); const probe = createGpuLeaseRegistry({ lockDirectory: directory }); try { await Effect.runPromise(primary.claim("llm", [proUuids[0], proUuids[1]])); await Effect.runPromise(blocker.claim("speech", [rtxUuid])); try { await Effect.runPromise(primary.replace("llm", [proUuids[2], rtxUuid])); throw new Error("expected a replacement GPU lease conflict"); } catch (error) { expect(error).toBeInstanceOf(GpuLeaseConflict); } expect(await Effect.runPromise(primary.snapshot())).toEqual([ { uuid: proUuids[0], owner: "llm" }, { uuid: proUuids[1], owner: "llm" }, ]); await Effect.runPromise(probe.claim("llm", [proUuids[2]])); await Effect.runPromise(probe.release("llm")); await Effect.runPromise(blocker.release("speech")); await Effect.runPromise(primary.replace("llm", [proUuids[2], rtxUuid])); expect(await Effect.runPromise(probe.claim("speech", [proUuids[0], proUuids[1]]))).toEqual([ { uuid: proUuids[0], owner: "speech" }, { uuid: proUuids[1], owner: "speech" }, ]); } finally { await Effect.runPromise(primary.release("llm")); await Effect.runPromise(blocker.release("speech")); await Effect.runPromise(probe.release("llm")); await Effect.runPromise(probe.release("speech")); await rm(directory, { recursive: true, force: true }); }});
import { expect, test } from "bun:test";
import { mkdir, mkdtemp, rm, utimes, writeFile } from "node:fs/promises";
import { tmpdir } from "node:os";
import { join } from "node:path";
import { Effect } from "effect";
import { asRecipeId, type GpuInfo, type Recipe } from "../models/types";
import {
createGpuLeaseRegistry,
GpuLeaseConflict,
GpuLeaseLockFailure,
resolveRecipeGpuUuids,
type GpuLease,
} from "./gpu-leases";
const proUuids = [
"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 rtxUuid = "GPU-00000000-0000-0000-0000-000000003090";
const temporaryLockDirectory = (): Promise<string> =>
mkdtemp(join(tmpdir(), "local-studio-gpu-leases-test-"));
const lockPath = (directory: string, uuid: string): string =>
join(directory, `${uuid.toLowerCase()}.lock`);
const writeHostLock = async (
directory: string,
uuid: string,
pid: number,
processStartToken: string | null,
): Promise<void> => {
await mkdir(directory, { recursive: true, mode: 0o700 });
await writeFile(
lockPath(directory, uuid),
JSON.stringify({
version: 1,
uuid,
owner: "speech",
pid,
processStartToken,
registryId: "stale-registry",
}),
{ mode: 0o600 },
);
};
function gpu(index: number, name: string, uuid: string): GpuInfo {
return {
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,
};
}
function recipe(
environmentVariables: Record<string, string> | null = null,
extraArguments: Record<string, unknown> = {},
): Recipe {
return {
id: asRecipeId("gpu-lease-test"),
name: "GPU lease test",
model_path: "model",
vision: null,
backend: "vllm",
runtime: { kind: "managed_venv", ref: "vllm-latest" },
env_vars: environmentVariables,
tensor_parallel_size: 1,
pipeline_parallel_size: 1,
max_model_len: 4096,
gpu_memory_utilization: 0.9,
kv_cache_dtype: "auto",
max_num_seqs: 1,
trust_remote_code: false,
tool_call_parser: null,
reasoning_parser: null,
enable_auto_tool_choice: false,
quantization: null,
dtype: null,
host: "127.0.0.1",
port: 8000,
served_model_name: null,
python_path: null,
extra_args: extraArguments,
max_thinking_tokens: null,
thinking_mode: "auto",
};
}
function fiveGpuHost(): GpuInfo[] {
return [
...proUuids.slice(0, 3).map((uuid, index) => gpu(index, "NVIDIA RTX PRO 6000 Blackwell", uuid)),
gpu(3, "NVIDIA GeForce RTX 3090", rtxUuid),
gpu(4, "NVIDIA RTX PRO 6000 Blackwell", proUuids[3]),
];
}
test("resolves the four PRO recipe without leasing the 3090", () => {
const resolution = resolveRecipeGpuUuids(
recipe({ CUDA_VISIBLE_DEVICES: "0,1,2,4" }),
fiveGpuHost(),
);
expect(resolution).toEqual({
source: "recipe",
selector: "0,1,2,4",
uuids: [...proUuids],
unresolvedTokens: [],
});
});
test("reports an atomic all-GPU lease conflict", async () => {
const registry = createGpuLeaseRegistry();
const allUuids = resolveRecipeGpuUuids(recipe(), fiveGpuHost()).uuids;
await Effect.runPromise(registry.claim("llm", proUuids));
await Effect.runPromise(registry.claim("speech", [rtxUuid]));
try {
await Effect.runPromise(registry.replace("llm", allUuids));
throw new Error("expected a GPU lease conflict");
} catch (error) {
expect(error).toBeInstanceOf(GpuLeaseConflict);
if (error instanceof GpuLeaseConflict) {
expect(error.requestedBy).toBe("llm");
expect(error.conflicts).toEqual([{ uuid: rtxUuid, heldBy: "speech" }]);
}
}
const unchanged: GpuLease[] = [
...proUuids.map((uuid): GpuLease => ({ uuid, owner: "llm" })),
{ uuid: rtxUuid, owner: "speech" },
];
expect(await Effect.runPromise(registry.snapshot())).toEqual(unchanged);
});
test("keeps UUID selectors stable when GPU indices reorder", () => {
const reordered = [
gpu(0, "NVIDIA RTX PRO 6000 Blackwell", proUuids[3]),
gpu(4, "NVIDIA RTX PRO 6000 Blackwell", proUuids[0]),
];
const selector = `${proUuids[0]},${proUuids[3]}`;
expect(
resolveRecipeGpuUuids(recipe({ CUDA_VISIBLE_DEVICES: selector }), reordered).uuids,
).toEqual([proUuids[0], proUuids[3]]);
});
test("reports unresolved aliased selectors without falling back", () => {
const resolution = resolveRecipeGpuUuids(
recipe(null, { "cuda-visible-devices": `${proUuids[1]},99` }),
fiveGpuHost(),
);
expect(resolution.uuids).toEqual([proUuids[1]]);
expect(resolution.unresolvedTokens).toEqual(["99"]);
});
test("replaces and releases only the requesting owner leases", async () => {
const registry = createGpuLeaseRegistry();
await Effect.runPromise(registry.claim("llm", [proUuids[0], proUuids[1]]));
await Effect.runPromise(registry.claim("speech", [rtxUuid]));
await Effect.runPromise(registry.replace("llm", [proUuids[2], proUuids[3]]));
await Effect.runPromise(registry.release("speech"));
await Effect.runPromise(registry.release("llm", [proUuids[2]]));
expect(await Effect.runPromise(registry.snapshot())).toEqual([
{ uuid: proUuids[3], owner: "llm" },
]);
});
test("conflicts across registries in the same live process", async () => {
const directory = await temporaryLockDirectory();
const first = createGpuLeaseRegistry({ lockDirectory: directory });
const second = createGpuLeaseRegistry({ lockDirectory: directory });
try {
await Effect.runPromise(first.claim("speech", [rtxUuid]));
expect(await Effect.runPromise(first.claim("speech", [rtxUuid]))).toEqual([
{ uuid: rtxUuid, owner: "speech" },
]);
try {
await Effect.runPromise(second.claim("speech", [rtxUuid]));
throw new Error("expected a cross-registry GPU lease conflict");
} catch (error) {
expect(error).toBeInstanceOf(GpuLeaseConflict);
if (error instanceof GpuLeaseConflict) {
expect(error.conflicts).toEqual([{ uuid: rtxUuid, heldBy: "speech" }]);
}
}
await Effect.runPromise(first.release("speech"));
expect(await Effect.runPromise(second.claim("speech", [rtxUuid]))).toEqual([
{ uuid: rtxUuid, owner: "speech" },
]);
} finally {
await Effect.runPromise(first.release("speech"));
await Effect.runPromise(second.release("speech"));
await rm(directory, { recursive: true, force: true });
}
});
test("reclaims a lock from a dead process", async () => {
const directory = await temporaryLockDirectory();
const registry = createGpuLeaseRegistry({ lockDirectory: directory });
try {
await writeHostLock(directory, rtxUuid, 2_147_483_647, "1");
expect(await Effect.runPromise(registry.claim("speech", [rtxUuid]))).toEqual([
{ uuid: rtxUuid, owner: "speech" },
]);
} finally {
await Effect.runPromise(registry.release("speech"));
await rm(directory, { recursive: true, force: true });
}
});
test("reclaims a reused Linux PID by process start token", async () => {
if (process.platform !== "linux") return;
const directory = await temporaryLockDirectory();
const registry = createGpuLeaseRegistry({ lockDirectory: directory });
try {
await writeHostLock(directory, rtxUuid, process.pid, "0");
expect(await Effect.runPromise(registry.claim("speech", [rtxUuid]))).toEqual([
{ uuid: rtxUuid, owner: "speech" },
]);
} finally {
await Effect.runPromise(registry.release("speech"));
await rm(directory, { recursive: true, force: true });
}
});
test("fails closed on an unverifiable host lock record", async () => {
const directory = await temporaryLockDirectory();
const registry = createGpuLeaseRegistry({ lockDirectory: directory });
try {
await writeFile(lockPath(directory, rtxUuid), "not-json", { mode: 0o600 });
await expect(Effect.runPromise(registry.claim("speech", [rtxUuid]))).rejects.toBeInstanceOf(
GpuLeaseLockFailure,
);
} finally {
await rm(directory, { recursive: true, force: true });
}
});
test("recovers a stale reaper left during dead-lock reclamation", async () => {
const directory = await temporaryLockDirectory();
const registry = createGpuLeaseRegistry({ lockDirectory: directory });
const reaper = `${lockPath(directory, rtxUuid)}.reaper`;
try {
await writeHostLock(directory, rtxUuid, 2_147_483_647, "1");
await mkdir(reaper, { mode: 0o700 });
const old = new Date(Date.now() - 10_000);
await utimes(reaper, old, old);
expect(await Effect.runPromise(registry.claim("speech", [rtxUuid]))).toEqual([
{ uuid: rtxUuid, owner: "speech" },
]);
} finally {
await Effect.runPromise(registry.release("speech"));
await rm(directory, { recursive: true, force: true });
}
});
test("rolls back new host locks before a failed replacement", async () => {
const directory = await temporaryLockDirectory();
const primary = createGpuLeaseRegistry({ lockDirectory: directory });
const blocker = createGpuLeaseRegistry({ lockDirectory: directory });
const probe = createGpuLeaseRegistry({ lockDirectory: directory });
try {
await Effect.runPromise(primary.claim("llm", [proUuids[0], proUuids[1]]));
await Effect.runPromise(blocker.claim("speech", [rtxUuid]));
try {
await Effect.runPromise(primary.replace("llm", [proUuids[2], rtxUuid]));
throw new Error("expected a replacement GPU lease conflict");
} catch (error) {
expect(error).toBeInstanceOf(GpuLeaseConflict);
}
expect(await Effect.runPromise(primary.snapshot())).toEqual([
{ uuid: proUuids[0], owner: "llm" },
{ uuid: proUuids[1], owner: "llm" },
]);
await Effect.runPromise(probe.claim("llm", [proUuids[2]]));
await Effect.runPromise(probe.release("llm"));
await Effect.runPromise(blocker.release("speech"));
await Effect.runPromise(primary.replace("llm", [proUuids[2], rtxUuid]));
expect(await Effect.runPromise(probe.claim("speech", [proUuids[0], proUuids[1]]))).toEqual([
{ uuid: proUuids[0], owner: "speech" },
{ uuid: proUuids[1], owner: "speech" },
]);
} finally {
await Effect.runPromise(primary.release("llm"));
await Effect.runPromise(blocker.release("speech"));
await Effect.runPromise(probe.release("llm"));
await Effect.runPromise(probe.release("speech"));
await rm(directory, { recursive: true, force: true });
}
});