196 lines · 6.0 KB · TypeScript
RawDownload
// CRITICALimport { existsSync } from "node:fs";import { resolve } from "node:path";import type { Config } from "../config/env";import type {  RuntimeBackendInfo,  RuntimeCudaInfo,  RuntimePlatformInfo,  RuntimePlatformKind,  RuntimeTorchBuildInfo,  SystemRuntimeInfo,} from "../types/models";import { getGpuInfo } from "./gpu";import { getVllmRuntimeInfo } from "./vllm-runtime";import { resolveBinary, runCommand } from "./command/command-utilities";import { getTorchBuildInfo } from "./torch-info";import { getRocmInfo, resolveRocmSmiTool } from "./rocm-info"; const extractCudaVersion = (output: string): string | null => {  const match = output.match(/CUDA Version\s*:\s*([0-9.]+)/i);  if (match) {    return match[1] ?? null;  }  return null;}; const extractNvccVersion = (output: string): string | null => {  const match = output.match(/release\s+([0-9.]+)/i);  if (match) {    return match[1] ?? null;  }  return null;}; export const getCudaInfo = (): RuntimeCudaInfo => {  const nvidiaSmi = process.env["NVIDIA_SMI_PATH"] || "nvidia-smi";  let driverVersion: string | null = null;  let cudaVersion: string | null = null;   const driverResult = runCommand(nvidiaSmi, ["--query-gpu=driver_version", "--format=csv,noheader,nounits"]);  if (driverResult.status === 0 && driverResult.stdout) {    driverVersion = driverResult.stdout.split("\n")[0]?.trim() || null;  }   const smiResult = runCommand(nvidiaSmi, []);  if (smiResult.status === 0) {    cudaVersion = extractCudaVersion(smiResult.stdout) ?? extractCudaVersion(smiResult.stderr);  }   if (!cudaVersion) {    const nvccResult = runCommand("nvcc", ["--version"]);    if (nvccResult.status === 0) {      cudaVersion = extractNvccVersion(nvccResult.stdout) ?? extractNvccVersion(nvccResult.stderr);    }  }   return {    driver_version: driverVersion,    cuda_version: cudaVersion,  };}; const detectPlatformKind = (args: {  forcedSmiTool: string | undefined;  torch: RuntimeTorchBuildInfo;  hasNvidiaSmi: boolean;  hasRocmSmi: boolean;}): RuntimePlatformKind => {  const forced = args.forcedSmiTool?.trim();  if (forced === "nvidia-smi") return "cuda";  if (forced === "amd-smi" || forced === "rocm-smi") return "rocm";   if (args.torch.torch_hip) return "rocm";  if (args.torch.torch_cuda) return "cuda";   if (args.hasNvidiaSmi) return "cuda";  if (args.hasRocmSmi) return "rocm";  return "unknown";}; const getSglangRuntimeInfo = (config: Config): RuntimeBackendInfo => {  const python = config.sglang_python || "python3";  const result = runCommand(python, [    "-c",    "import json, sys\ntry:\n import sglang\n print(json.dumps({'version': getattr(sglang, '__version__', None), 'python': sys.executable}))\nexcept Exception:\n print(json.dumps({'version': None, 'python': sys.executable}))",  ]);   if (result.status !== 0) {    return {      installed: false,      version: null,      python_path: config.sglang_python ?? null,    };  }   let parsed: { version?: string | null; python?: string | null } | null = null;  try {    parsed = JSON.parse(result.stdout) as { version?: string | null; python?: string | null };  } catch {    parsed = null;  }   return {    installed: Boolean(parsed?.version),    version: parsed?.version ?? null,    python_path: parsed?.python ?? config.sglang_python ?? null,  };}; const parseLlamaVersion = (output: string): string | null => {  if (!output) return null;  const match = output.match(/version\s*[:=]\s*([^\s]+)/i);  if (match) {    return match[1] ?? null;  }  const fallback = output.split("\n")[0]?.trim();  return fallback || null;}; const getLlamacppRuntimeInfo = (config: Config): RuntimeBackendInfo => {  const configured = config.llama_bin || "llama-server";  const resolved = resolveBinary(configured) ?? (existsSync(configured) ? resolve(configured) : null);  const binary = resolved ?? configured;   const versionResult = runCommand(binary, ["--version"]);  if (versionResult.status !== 0) {    const helpResult = runCommand(binary, ["--help"]);    if (helpResult.status !== 0) {      return {        installed: false,        version: null,        binary_path: resolved,      };    }    const version = parseLlamaVersion(helpResult.stdout) ?? parseLlamaVersion(helpResult.stderr);    return {      installed: Boolean(version),      version,      binary_path: resolved,    };  }   const version = parseLlamaVersion(versionResult.stdout) ?? parseLlamaVersion(versionResult.stderr);  return {    installed: Boolean(version),    version,    binary_path: resolved,  };}; export const getSystemRuntimeInfo = async (config: Config): Promise<SystemRuntimeInfo> => {  const gpus = getGpuInfo();  const types = Array.from(new Set(gpus.map((gpu) => gpu.name).filter((name) => name && name !== "Unknown")));   const [vllmInfo, sglangInfo] = await Promise.all([    getVllmRuntimeInfo(),    Promise.resolve(getSglangRuntimeInfo(config)),  ]);   const llamaInfo = getLlamacppRuntimeInfo(config);   const pythonForTorch = config.sglang_python || vllmInfo.python_path || "python3";  const torch = getTorchBuildInfo(pythonForTorch);   const forcedSmiTool = process.env["VLLM_STUDIO_GPU_SMI_TOOL"];  const hasNvidiaSmi = Boolean(resolveBinary(process.env["NVIDIA_SMI_PATH"] || "nvidia-smi"));  const rocmSmiTool = resolveRocmSmiTool();  const hasRocmSmi = Boolean(rocmSmiTool);  const kind = detectPlatformKind({ forcedSmiTool, torch, hasNvidiaSmi, hasRocmSmi });   const platform: RuntimePlatformInfo = {    kind,    vendor: kind === "cuda" ? "nvidia" : kind === "rocm" ? "amd" : null,    rocm: kind === "rocm" ? getRocmInfo(rocmSmiTool) : null,    torch,  };   return {    platform,    cuda: kind === "cuda" ? getCudaInfo() : { driver_version: null, cuda_version: null },    gpus: {      count: gpus.length,      types,    },    backends: {      vllm: {        installed: vllmInfo.installed,        version: vllmInfo.version,        python_path: vllmInfo.python_path,        binary_path: vllmInfo.vllm_bin,      },      sglang: sglangInfo,      llamacpp: llamaInfo,    },  };};