488 lines · 15.8 KB · TypeScript
RawDownload
import { afterAll, describe, expect, test } from "bun:test";import { mkdtempSync, rmSync } from "node:fs";import { tmpdir } from "node:os";import { join } from "node:path";import { Cause, Effect } from "effect";import type {  HandleReference,  HostProfile,  InstanceRecord,  LaunchFailure,  ServingOptions,} from "../src/modules/compute/contracts";import { toEvent, toHttp } from "../src/modules/compute/failures";import { makeInstanceStore } from "../src/modules/compute/instances/store";import type { Launcher } from "../src/modules/compute/launchers/launcher";import {  makeComputeService,  type ComputeLaunchInput,  type ComputeService,} from "../src/modules/compute/lifecycle"; const run = <A, E>(effect: Effect.Effect<A, E>): Promise<A> => Effect.runPromise(effect);const runExit = <A, E>(effect: Effect.Effect<A, E>) => Effect.runPromiseExit(effect); const roots: string[] = [];const freshRoot = (): string => {  const root = mkdtempSync(join(tmpdir(), "compute-test-"));  roots.push(root);  return root;};afterAll(() => {  for (const root of roots) rmSync(root, { recursive: true, force: true });}); const linuxHost: HostProfile = {  nodeId: "self",  platform: "linux",  arch: "x64",  accelerator: "cuda",  unifiedMemory: false,  wsl: false,  docker: false,  dockerGpu: false,  deviceCount: 2,}; const options: ServingOptions = {  tensorParallel: 1,  pipelineParallel: 1,  maxContextLength: 4096,  memoryFraction: 0.9,  maxConcurrentRequests: 32,  kvCacheDtype: null,  dtype: null,  quantization: null,  trustRemoteCode: false,  toolCallParser: null,  reasoningParser: null,}; const input = (overrides: Partial<ComputeLaunchInput> = {}): ComputeLaunchInput => ({  name: "test-model",  engine: "llamacpp",  recipeId: "r1",  runtime: "process",  deviceCount: 1,  modelPath: "/models/test.gguf",  servedModelName: "test",  options,  extraArgs: [],  env: {},  binary: "llama-server",  ...overrides,}); /** Launcher whose world is a mutable script: which handles are alive, what start does. */interface FakeWorld {  alivePids: Set<number>;  started: InstanceRecord[];  stopped: HandleReference[];  nextPid: number;  startFailure: LaunchFailure | null;  healthyPorts: Set<number>;} const fakeLauncher = (world: FakeWorld): Launcher => ({  start: (_plan, record) => {    if (world.startFailure) return Effect.fail(world.startFailure);    const pid = world.nextPid++;    world.alivePids.add(pid);    world.started.push(record);    return Effect.succeed({ kind: "process", pid, startToken: null });  },  alive: (ref) => Effect.succeed(ref.kind === "process" && world.alivePids.has(ref.pid)),  owns: (ref) => Effect.succeed(ref.kind === "process" && world.alivePids.has(ref.pid)),  stop: (ref) =>    Effect.sync(() => {      if (ref.kind === "process") world.alivePids.delete(ref.pid);      world.stopped.push(ref);    }),  logTail: () => Effect.succeed("fake log tail"),}); const makeWorld = (): FakeWorld => ({  alivePids: new Set(),  started: [],  stopped: [],  nextPid: 1000,  startFailure: null,  healthyPorts: new Set(),}); const makeService = (  world: FakeWorld,  devices: readonly string[] = ["GPU-a", "GPU-b"],): { compute: ComputeService; store: ReturnType<typeof makeInstanceStore>; events: string[] } => {  const store = makeInstanceStore(freshRoot());  const events: string[] = [];  const compute = makeComputeService({    store,    launcherFor: () => fakeLauncher(world),    host: () => Effect.succeed(linuxHost),    freeDevices: () => Effect.succeed(devices),    onEvent: (name, stage) => Effect.sync(() => void events.push(`${name}:${stage}`)),  });  return { compute, store, events };}; // The fake never serves HTTP, so "healthy" can only come from a real server. For the// ready-path test we stand up a real Bun server on the allocated port instead. describe("reserve is the lease", () => {  test("devices of a live record are held; dropping the record frees them", async () => {    const world = makeWorld();    const { store } = makeService(world);    const alive = () => Effect.succeed(true);    const first = await run(      store.reserve(        {          name: "a",          nodeId: "self",          engine: "llamacpp",          recipeId: "r",          runtime: "process",          candidates: ["GPU-a", "GPU-b"],          need: 2,          basePort: 8081,          readyDeadlineMs: 1000,          shareable: false,        },        alive,      ),    );    expect(first.devices).toEqual(["GPU-a", "GPU-b"]);     // Everything is now held — a second reservation must fail with no-capacity.    const second = await runExit(      store.reserve(        {          name: "b",          nodeId: "self",          engine: "llamacpp",          recipeId: "r",          runtime: "process",          candidates: ["GPU-a", "GPU-b"],          need: 1,          basePort: 8081,          readyDeadlineMs: 1000,          shareable: false,        },        alive,      ),    );    expect(second._tag).toBe("Failure");     // Drop the record: capacity comes back with no release call anywhere.    store.drop("a");    const third = await run(      store.reserve(        {          name: "c",          nodeId: "self",          engine: "llamacpp",          recipeId: "r",          runtime: "process",          candidates: ["GPU-a", "GPU-b"],          need: 1,          basePort: 8081,          readyDeadlineMs: 1000,          shareable: false,        },        alive,      ),    );    expect(third.devices).toEqual(["GPU-a"]);  });   test("a dead record's devices are not held", async () => {    const world = makeWorld();    const { store } = makeService(world);    // Simulate a record whose process died: alive() answers false.    const dead = () => Effect.succeed(false);    await run(      store.reserve(        {          name: "a",          nodeId: "self",          engine: "llamacpp",          recipeId: "r",          runtime: "process",          candidates: ["GPU-a"],          need: 1,          basePort: 8081,          readyDeadlineMs: 1000,          shareable: false,        },        dead,      ),    );    const record = store.read("a");    expect(record).not.toBeNull();    if (!record) return;    store.write({ ...record, ref: { kind: "process", pid: 999999, startToken: null } });    const held = await run(store.heldDevices(dead));    expect(held.size).toBe(0);  });   test("concurrent reservations under the placement lock never double-book", async () => {    const world = makeWorld();    const { store } = makeService(world);    const alive = () => Effect.succeed(true);    const reserveOne = (name: string) =>      store.reserve(        {          name,          nodeId: "self",          engine: "llamacpp",          recipeId: "r",          runtime: "process",          candidates: ["GPU-a", "GPU-b"],          need: 1,          basePort: 8081,          readyDeadlineMs: 1000,          shareable: false,        },        alive,      );    const results = await run(      Effect.all([reserveOne("one"), reserveOne("two")], { concurrency: "unbounded" }).pipe(        Effect.map((records) => records.map((record) => record.devices).flat()),      ),    );    // Two reservations, two distinct devices — never the same one twice.    expect(new Set(results).size).toBe(2);  });   test("ports allocate upward from the engine base without collision", async () => {    const world = makeWorld();    const { store } = makeService(world);    const alive = () => Effect.succeed(true);    // High base so real listeners on this machine cannot interfere with the assertion.    const base = 42_081;    const reserveNamed = (name: string) =>      run(        store.reserve(          {            name,            nodeId: "self",            engine: "llamacpp",            recipeId: "r",            runtime: "process",            candidates: [`GPU-${name}`],            need: 1,            basePort: base,            readyDeadlineMs: 1000,            shareable: false,          },          alive,        ),      );    const one = await reserveNamed("one");    const two = await reserveNamed("two");    expect(one.port).toBeGreaterThanOrEqual(base);    expect(two.port).toBeGreaterThan(one.port);  });   test("allocatePort refuses a port an unrelated process already holds", () => {    const world = makeWorld();    const { store } = makeService(world);    const squatter = Bun.serve({ port: 0, fetch: () => new Response("ok") });    try {      const taken = squatter.port as number;      expect(store.allocatePort(taken)).toBeGreaterThan(taken);    } finally {      squatter.stop(true);    }  });}); describe("launch failure paths", () => {  test("unsupported engine/host fails before touching devices", async () => {    const world = makeWorld();    const { compute, store } = makeService(world);    const exit = await runExit(compute.launch(input({ engine: "mlx" })));    expect(exit._tag).toBe("Failure");    expect(store.all()).toHaveLength(0);    expect(world.started).toHaveLength(0);  });   test("spawn failure drops the record and frees the devices", async () => {    const world = makeWorld();    world.startFailure = { kind: "spawn-failed", detail: "binary not found" };    const { compute, store, events } = makeService(world);    const exit = await runExit(compute.launch(input()));    expect(exit._tag).toBe("Failure");    // The reservation must not leak: record gone, devices free.    expect(store.all()).toHaveLength(0);    expect(events).toContain("test-model:error");  });   test("exited-early: daemon dies while waiting -> record dropped, log tail captured", async () => {    const world = makeWorld();    const { compute, store } = makeService(world);    // Kill the "process" the moment it starts.    const original = fakeLauncher(world).start;    void original;    const launchEffect = compute.launch(input());    const exitPromise = runExit(launchEffect);    // Give the launch a moment to spawn, then kill the pid.    await new Promise((resolve) => setTimeout(resolve, 150));    world.alivePids.clear();    const exit = await exitPromise;    expect(exit._tag).toBe("Failure");    expect(store.all()).toHaveLength(0);  });   test("cancel mid-wait -> cancelled failure, record dropped", async () => {    const world = makeWorld();    const { compute, store } = makeService(world);    const exitPromise = runExit(compute.launch(input()));    await new Promise((resolve) => setTimeout(resolve, 150));    await run(compute.cancel("test-model"));    const exit = await exitPromise;    expect(exit._tag).toBe("Failure");    if (exit._tag === "Failure") {      const failure = Cause.findErrorOption(exit.cause);      expect(failure._tag).toBe("Some");      if (failure._tag === "Some") {        expect((failure.value as LaunchFailure).kind).toBe("cancelled");      }    }    expect(store.all()).toHaveLength(0);  });   test("second launch under the same name while starting -> already-running", async () => {    const world = makeWorld();    const { compute } = makeService(world);    const firstPromise = runExit(compute.launch(input()));    await new Promise((resolve) => setTimeout(resolve, 150));    const second = await runExit(compute.launch(input()));    expect(second._tag).toBe("Failure");    if (second._tag === "Failure") {      const failure = Cause.findErrorOption(second.cause);      if (failure._tag === "Some") {        expect((failure.value as LaunchFailure).kind).toBe("already-running");      }    }    await run(compute.cancel("test-model"));    await firstPromise;  });}); describe("ready path with a real health endpoint", () => {  test("launch succeeds once /health answers 200", async () => {    const world = makeWorld();    const { compute, store, events } = makeService(world);    // The port is allocated inside launch; discover it from the reservation record    // (written before spawn) and only then stand up the health endpoint on it.    const launchPromise = run(compute.launch(input()));    let port: number | null = null;    for (let attempt = 0; attempt < 100 && port === null; attempt += 1) {      port = store.read("test-model")?.port ?? null;      if (port === null) await new Promise((resolve) => setTimeout(resolve, 20));    }    expect(port).not.toBeNull();    if (port === null) return;    const server = Bun.serve({      port,      fetch: (request) =>        new URL(request.url).pathname === "/health"          ? new Response("ok")          : new Response("no", { status: 404 }),    });    try {      const record = await launchPromise;      expect(record.ref?.kind).toBe("process");      expect(record.port).toBe(port);      expect(events).toContain("test-model:ready");      const views = await run(compute.instances());      expect(views).toHaveLength(1);      expect(views[0]?.state).toBe("ready");      // stop() tears down and frees.      expect(await run(compute.stop("test-model"))).toBe(true);      expect(store.all()).toHaveLength(0);      expect(world.stopped).toHaveLength(1);    } finally {      server.stop(true);    }  });}); describe("supervisor", () => {  test("reaps records whose handle is gone; leaves live ones alone", async () => {    const world = makeWorld();    const { compute, store } = makeService(world);    const liveRef: HandleReference = { kind: "process", pid: 2000, startToken: null };    world.alivePids.add(2000);    const base: InstanceRecord = {      name: "live",      nodeId: "self",      engine: "llamacpp",      recipeId: "r",      runtime: "process",      ref: liveRef,      port: 8081,      devices: ["GPU-a"],      nonce: "n1",      startedAt: new Date().toISOString(),      readyDeadlineAt: new Date(Date.now() + 60_000).toISOString(),    };    store.write(base);    store.write({      ...base,      name: "dead",      ref: { kind: "process", pid: 3000, startToken: null },      port: 8082,      devices: ["GPU-b"],    });    const reaped = await run(compute.superviseOnce());    expect(reaped).toBe(1);    expect(store.read("live")).not.toBeNull();    expect(store.read("dead")).toBeNull();  });   test("stale reservations (no handle) are reaped after the grace window", async () => {    const world = makeWorld();    const { compute, store } = makeService(world);    store.write({      name: "crashed-before-spawn",      nodeId: "self",      engine: "llamacpp",      recipeId: "r",      runtime: "process",      ref: null,      port: 8081,      devices: ["GPU-a"],      nonce: "n",      startedAt: new Date(Date.now() - 120_000).toISOString(),      readyDeadlineAt: new Date(Date.now() - 60_000).toISOString(),    });    expect(await run(compute.superviseOnce())).toBe(1);    expect(store.all()).toHaveLength(0);  });}); describe("failure mappers", () => {  test("every failure kind maps to exactly one status and one stage", () => {    const failures: LaunchFailure[] = [      { kind: "unsupported", engine: "vllm", reason: "no Metal backend" },      { kind: "already-running", name: "x" },      { kind: "no-capacity", need: 2, free: 0 },      { kind: "install-failed", engine: "vllm", detail: "pip failed" },      { kind: "spawn-failed", detail: "ENOENT" },      { kind: "exited-early", exitCode: 1, signal: null, logTail: "CUDA OOM" },      { kind: "unhealthy-timeout", waitedMs: 60_000, logTail: "loading..." },      { kind: "cancelled" },    ];    const statuses = failures.map((failure) => toHttp(failure).status);    expect(statuses).toEqual([422, 409, 409, 503, 503, 503, 503, 400]);    // Cancellation is a stage, not a substring match on the message.    expect(toEvent({ kind: "cancelled" }).stage).toBe("cancelled");    expect(toEvent({ kind: "spawn-failed", detail: "operation cancelled by user" }).stage).toBe(      "error",    );    // The crash log rides along at a single truncation length.    expect(toHttp(failures[5] as LaunchFailure).detail).toContain("CUDA OOM");  });});

Keyboard shortcuts

General

Show keyboard shortcuts
?
Focus the current search
/

Issues and pull requests

Open filters
f

Issues, pull requests and Inbox

Select next item
j↓
Select previous item
k↑
Open selected item
Enter
Dismiss filters or clear selection
Esc

Repository code

Go to file
t