import { runCommandAsync, type AsyncCommandResult } from "../../../core/command";import { resolveNvidiaSmiBinary } from "./smi-tools"; const FULL_NVIDIA_UUID = /^GPU-[0-9a-fA-F]{8}-[0-9a-fA-F]{4}-[0-9a-fA-F]{4}-[0-9a-fA-F]{4}-[0-9a-fA-F]{12}$/;const QUERY_ARGS = ["--query-compute-apps=gpu_uuid,pid", "--format=csv,noheader,nounits"] as const; export interface NvidiaComputeProcessDependencies { readonly resolveBinary: () => string | null; readonly runCommand: ( command: string, args: string[], ) => Promise<Pick<AsyncCommandResult, "exitConfirmed" | "status" | "stdout">>;} const dependencies: NvidiaComputeProcessDependencies = { resolveBinary: resolveNvidiaSmiBinary, runCommand: (command, args) => runCommandAsync(command, args, { timeoutMs: 5_000, maxOutputBytes: 256 * 1024 }),}; const canonicalUuid = (uuid: string): string => `GPU-${uuid.slice(4).toLowerCase()}`; const computeGpuUuids = (stdout: string): readonly string[] => { const uuids = new Set<string>(); for (const line of stdout .split("\n") .map((value) => value.trim()) .filter(Boolean)) { const [uuid, pid, ...extra] = line.split(",").map((value) => value.trim()); if (!uuid || !pid || extra.length > 0 || !FULL_NVIDIA_UUID.test(uuid) || !/^\d+$/.test(pid)) { throw new Error("NVIDIA compute process output is invalid"); } uuids.add(canonicalUuid(uuid)); } return [...uuids];}; export const queryNvidiaComputeGpuUuids = async ( injected: NvidiaComputeProcessDependencies = dependencies,): Promise<readonly string[]> => { const binary = injected.resolveBinary(); if (!binary) throw new Error("NVIDIA compute process telemetry is unavailable"); const result = await injected.runCommand(binary, [...QUERY_ARGS]); if (result.status !== 0 || result.exitConfirmed === false) { throw new Error("NVIDIA compute process telemetry failed"); } return computeGpuUuids(result.stdout);};
import { runCommandAsync, type AsyncCommandResult } from "../../../core/command";
import { resolveNvidiaSmiBinary } from "./smi-tools";
const FULL_NVIDIA_UUID =
/^GPU-[0-9a-fA-F]{8}-[0-9a-fA-F]{4}-[0-9a-fA-F]{4}-[0-9a-fA-F]{4}-[0-9a-fA-F]{12}$/;
const QUERY_ARGS = ["--query-compute-apps=gpu_uuid,pid", "--format=csv,noheader,nounits"] as const;
export interface NvidiaComputeProcessDependencies {
readonly resolveBinary: () => string | null;
readonly runCommand: (
command: string,
args: string[],
) => Promise<Pick<AsyncCommandResult, "exitConfirmed" | "status" | "stdout">>;
}
const dependencies: NvidiaComputeProcessDependencies = {
resolveBinary: resolveNvidiaSmiBinary,
runCommand: (command, args) =>
runCommandAsync(command, args, { timeoutMs: 5_000, maxOutputBytes: 256 * 1024 }),
};
const canonicalUuid = (uuid: string): string => `GPU-${uuid.slice(4).toLowerCase()}`;
const computeGpuUuids = (stdout: string): readonly string[] => {
const uuids = new Set<string>();
for (const line of stdout
.split("\n")
.map((value) => value.trim())
.filter(Boolean)) {
const [uuid, pid, ...extra] = line.split(",").map((value) => value.trim());
if (!uuid || !pid || extra.length > 0 || !FULL_NVIDIA_UUID.test(uuid) || !/^\d+$/.test(pid)) {
throw new Error("NVIDIA compute process output is invalid");
}
uuids.add(canonicalUuid(uuid));
}
return [...uuids];
};
export const queryNvidiaComputeGpuUuids = async (
injected: NvidiaComputeProcessDependencies = dependencies,
): Promise<readonly string[]> => {
const binary = injected.resolveBinary();
if (!binary) throw new Error("NVIDIA compute process telemetry is unavailable");
const result = await injected.runCommand(binary, [...QUERY_ARGS]);
if (result.status !== 0 || result.exitConfirmed === false) {
throw new Error("NVIDIA compute process telemetry failed");
}
return computeGpuUuids(result.stdout);
};