290 lines · 8.8 KB · TypeScript
RawDownload
import { describe, expect, test } from "bun:test";import { chmodSync, mkdirSync, writeFileSync } from "node:fs";import { join } from "node:path";import { createConfig } from "../../../controller/src/config/env";import { createLogger } from "../../../controller/src/core/logger";import { createProcessManager } from "../../../controller/src/modules/engines/process/process-manager";import type { Recipe } from "../../../controller/src/modules/models/types";import { FakeProcessRunner } from "../../support/controller/fake-process-runner";import { registerControllerTestLifecycle, tempDir } from "./fixtures"; registerControllerTestLifecycle(); const LAUNCH_TIMEOUT_MS = 10_000; const testVllmBinary = (): string => {  const binary = join(tempDir, "bin", "vllm");  mkdirSync(join(tempDir, "bin"), { recursive: true });  writeFileSync(binary, "");  chmodSync(binary, 0o755);  return binary;}; const vllmRecipe = (  id: string,  runtime: Recipe["runtime"] = { kind: "system", ref: testVllmBinary() },  extraArguments: Record<string, unknown> = {},): Recipe => ({  id,  name: id,  model_path: join(tempDir, "models", "test-model"),  backend: "vllm",  runtime,  env_vars: null,  tensor_parallel_size: 1,  pipeline_parallel_size: 1,  max_model_len: 4096,  gpu_memory_utilization: 0.8,  kv_cache_dtype: "auto",  max_num_seqs: 8,  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: "test-model",  python_path: null,  extra_args: extraArguments,  max_thinking_tokens: null,  thinking_mode: null,}); const createManager = (runner: FakeProcessRunner) => {  const config = createConfig();  return {    config,    manager: createProcessManager(config, createLogger("error"), undefined, runner),  };}; const onlySpawn = (runner: FakeProcessRunner) => {  const spawns = runner.spawns();  expect(spawns).toHaveLength(1);  const spawn = spawns[0];  if (!spawn) throw new Error("Expected one process spawn");  return spawn;}; describe("process manager launch/stop via the process seam", () => {  test(    "launchModel spawns the selected vLLM runtime argv, reports the pid, and stops",    async () => {      const runner = new FakeProcessRunner();      const { config, manager } = createManager(runner);      const recipe = vllmRecipe("vllm-argv");       const result = await manager.launchModel(recipe);       expect(result.success).toBe(true);      expect(result.pid).toBe(42_001);      expect(result.log_file).toBe(join(config.data_dir, "logs", "vllm_vllm-argv.log"));       const spawn = onlySpawn(runner);      expect(spawn.command.split("/").at(-1)).toBe("vllm");      expect(spawn.args).toEqual([        "serve",        recipe.model_path,        "--host",        "127.0.0.1",        "--port",        String(config.inference_port),        "--served-model-name",        "test-model",        "--max-model-len",        "4096",        "--gpu-memory-utilization",        "0.8",        "--max-num-seqs",        "8",      ]);      expect(        runner.invocations.some(          (invocation) => invocation.kind === "runSync" && invocation.command === "ps",        ),      ).toBe(true);      if (result.pid === null) throw new Error("Expected launched process pid");      await expect(manager.killProcess(result.pid, false)).resolves.toBe(true);    },    LAUNCH_TIMEOUT_MS,  );   test(    "docker runtime removes the stale container and spawns the exact docker argv",    async () => {      const runner = new FakeProcessRunner();      const { config, manager } = createManager(runner);      const image = "vllm/vllm-custom:test";      const recipe = vllmRecipe("vllm-docker", { kind: "docker", ref: image });       const result = await manager.launchModel(recipe);      expect(result.success).toBe(true);      expect(runner.invocations).toContainEqual({        kind: "runSync",        command: "docker",        args: ["rm", "-f", "local-studio-vllm-docker"],      });       const spawn = onlySpawn(runner);      expect(spawn.command).toBe("docker");      expect(spawn.args).toEqual([        "run",        "--rm",        "--name",        "local-studio-vllm-docker",        "--gpus",        "all",        "--network",        "host",        "--ipc",        "host",        "--shm-size",        "32g",        "--ulimit",        "memlock=-1",        "--ulimit",        "stack=67108864",        "-e",        "XDG_CACHE_HOME=/cache/jit",        "-e",        "CUDA_CACHE_PATH=/cache/jit",        "-e",        "VLLM_CACHE_DIR=/cache/jit/vllm",        "-e",        "TRITON_CACHE_DIR=/cache/jit/triton",        "-v",        `${recipe.model_path}:${recipe.model_path}:ro`,        "-v",        "local-studio-jit-vllm-docker:/cache/jit",        image,        "/opt/venv/bin/vllm",        "serve",        recipe.model_path,        "--host",        "127.0.0.1",        "--port",        String(config.inference_port),        "--served-model-name",        "test-model",        "--max-model-len",        "4096",        "--gpu-memory-utilization",        "0.8",        "--max-num-seqs",        "8",      ]);    },    LAUNCH_TIMEOUT_MS,  );   test(    "docker runtime receives the coordinator's exact canonical UUID set",    async () => {      const runner = new FakeProcessRunner();      const { manager } = createManager(runner);      const image = "vllm/vllm-custom:test";      const recipe = vllmRecipe(        "vllm-docker-uuids",        { kind: "docker", ref: image },        { visible_devices: "3", env_vars: { CUDA_VISIBLE_DEVICES: "3" } },      );      const gpuUuids = [        "GPU-00000000-0000-0000-0000-000000000001",        "GPU-00000000-0000-0000-0000-000000000002",        "GPU-00000000-0000-0000-0000-000000000003",        "GPU-00000000-0000-0000-0000-000000000004",      ];       const result = await manager.launchModel(recipe, { gpuUuids });       expect(result.success).toBe(true);      const selector = gpuUuids.join(",");      const spawn = onlySpawn(runner);      expect(spawn.args.slice(4, 8)).toEqual([        "--gpus",        `"device=${selector}"`,        "-e",        `CUDA_VISIBLE_DEVICES=${selector}`,      ]);    },    LAUNCH_TIMEOUT_MS,  );   test("managed GPU selection rejects a custom launch command", async () => {    const runner = new FakeProcessRunner();    const { manager } = createManager(runner);    const key = "LOCAL_STUDIO_ALLOW_CUSTOM_LAUNCH_COMMAND";    const previous = process.env[key];    process.env[key] = "true";    try {      const result = await manager.launchModel(        vllmRecipe("vllm-custom-gpu", undefined, {          launch_command: "docker run --gpus all unsafe-image",        }),        { gpuUuids: ["GPU-00000000-0000-0000-0000-000000000001"] },      );       expect(result.success).toBe(false);      expect(result.message).toBe("Custom launch commands cannot use managed GPU selection");      expect(runner.spawns()).toEqual([]);    } finally {      if (previous === undefined) delete process.env[key];      else process.env[key] = previous;    }  });   test("stop confirmation cleans and retains known orphan workers", async () => {    const pid = 2_147_483_000;    const runner = new FakeProcessRunner().onRunSync(      (command, args) => command === "ps" && args.includes("pid=,ppid=,stat=,args="),      { stdout: `${pid} 1 S VLLM::Worker` },    );    const { manager } = createManager(runner);     await expect(manager.confirmInferenceStopped(8000)).resolves.toBe(false);    expect(runner.invocations).toContainEqual({      kind: "runSync",      command: "sudo",      args: ["-n", "kill", "-SIGTERM", String(pid)],    });  });   test(    "a fast-failing launch surfaces the exit code and captured output",    async () => {      const runner = new FakeProcessRunner().onSpawn("vllm", {        exitCode: 2,        exitAfterMs: 50,        stderrLines: ["usage: vllm serve", "error: unrecognized arguments: --bogus-flag"],      });      const { manager } = createManager(runner);       const result = await manager.launchModel(vllmRecipe("vllm-early-exit"));       expect(result.success).toBe(false);      expect(result.pid).toBeNull();      expect(result.message).toContain("Process exited early (code 2)");      expect(result.message).toContain("unrecognized arguments: --bogus-flag");    },    LAUNCH_TIMEOUT_MS,  );   test(    "a spawn error is surfaced as a failed launch",    async () => {      const runner = new FakeProcessRunner().onSpawn("vllm", {        spawnError: "spawn vllm ENOENT",      });      const { manager } = createManager(runner);       const result = await manager.launchModel(vllmRecipe("vllm-enoent"));       expect(result.success).toBe(false);      expect(result.pid).toBeNull();      expect(result.message).toContain("ENOENT");    },    LAUNCH_TIMEOUT_MS,  );});