import { spawnSync } from "node:child_process";import { dirname } from "node:path";import type { Recipe } from "../../models/types";import type { Backend } from "@local-studio/contracts/recipes";import { detectEngineFromArguments } from "../engine-spec";import { extractFlag as extractFlagUtility } from "../argument-utilities";import { isManagedPythonBackend, managedVenvPython } from "../runtimes/managed-venv";import type { Config } from "../../../config/env"; export { extractFlagUtility as extractFlag }; const splitCommand = (command: string): string[] => { const matches = command.match(/(?:[^\s"]+|"[^"]*")+/g) ?? []; return matches.map((token) => token.replace(/^"|"$/g, ""));}; export const detectBackend = (args: string[]): Backend | null => { if (args.length === 0) return null; return detectEngineFromArguments(args);}; export const listProcesses = (): Array<{ pid: number; args: string[] }> => { try { const result = spawnSync("ps", ["-eo", "pid=,args="]); if (result.status !== 0) { return []; } const output = result.stdout.toString("utf-8").trim(); if (!output) { return []; } return output .split("\n") .map((line) => { const trimmed = line.trim(); const match = trimmed.match(/^(\d+)\s+(.*)$/); if (!match) { return { pid: 0, args: [] }; } const pid = Number(match[1]); const args = splitCommand(match[2] ?? ""); return { pid, args }; }) .filter((entry) => entry.pid > 0 && entry.args.length > 0); } catch { return []; }}; export const buildEnvironment = ( recipe: Recipe, config?: Pick<Config, "data_dir">,): Record<string, string> => { const env: Record<string, string> = { ...process.env } as Record<string, string>; env["FLASHINFER_DISABLE_VERSION_CHECK"] = "1"; const venvBin = resolveVenvBinForRecipe(recipe, config?.data_dir); if (venvBin) { env["PATH"] = `${venvBin}:${env["PATH"] ?? ""}`; } const environmentVariables: Record<string, string> = {}; if (recipe.env_vars && typeof recipe.env_vars === "object") { for (const [key, value] of Object.entries(recipe.env_vars)) { if (value !== undefined && value !== null) { environmentVariables[String(key)] = String(value); } } } const extraEnvironment = recipe.extra_args["env_vars"] || recipe.extra_args["env-vars"] || recipe.extra_args["envVars"]; if (extraEnvironment && typeof extraEnvironment === "object") { for (const [key, value] of Object.entries(extraEnvironment as Record<string, unknown>)) { if (value !== undefined && value !== null) { environmentVariables[String(key)] = String(value); } } } for (const [key, value] of Object.entries(environmentVariables)) { env[key] = value; } const readExtraArgument = (key: string): unknown => { if (Object.prototype.hasOwnProperty.call(recipe.extra_args, key)) { return recipe.extra_args[key]; } const kebab = key.replace(/_/g, "-"); if (Object.prototype.hasOwnProperty.call(recipe.extra_args, kebab)) { return recipe.extra_args[kebab]; } const snake = key.replace(/-/g, "_"); if (Object.prototype.hasOwnProperty.call(recipe.extra_args, snake)) { return recipe.extra_args[snake]; } return undefined; }; const isDefined = (value: unknown): boolean => { return value !== undefined && value !== null && value !== false; }; const visibleDevices = readExtraArgument("visible_devices") ?? readExtraArgument("VISIBLE_DEVICES") ?? readExtraArgument("CUDA_VISIBLE_DEVICES") ?? readExtraArgument("cuda_visible_devices") ?? readExtraArgument("cuda-visible-devices"); const hipVisibleDevices = readExtraArgument("hip_visible_devices") ?? readExtraArgument("HIP_VISIBLE_DEVICES"); const rocrVisibleDevices = readExtraArgument("rocr_visible_devices") ?? readExtraArgument("ROCR_VISIBLE_DEVICES"); const forcedTool = (process.env["LOCAL_STUDIO_GPU_SMI_TOOL"] ?? "").trim().toLowerCase(); const platform = forcedTool === "nvidia-smi" ? "cuda" : forcedTool === "amd-smi" || forcedTool === "rocm-smi" ? "rocm" : "unknown"; if (isDefined(visibleDevices)) { const value = String(visibleDevices); if (platform === "cuda") { env["CUDA_VISIBLE_DEVICES"] = value; } else if (platform === "rocm") { env["HIP_VISIBLE_DEVICES"] = value; env["ROCR_VISIBLE_DEVICES"] = value; } else { env["CUDA_VISIBLE_DEVICES"] = value; env["HIP_VISIBLE_DEVICES"] = value; env["ROCR_VISIBLE_DEVICES"] = value; } } if (isDefined(hipVisibleDevices)) { env["HIP_VISIBLE_DEVICES"] = String(hipVisibleDevices); } if (isDefined(rocrVisibleDevices)) { env["ROCR_VISIBLE_DEVICES"] = String(rocrVisibleDevices); } return env;}; function resolveVenvBinForRecipe(recipe: Recipe, dataDirectory?: string): string | null { if ( recipe.runtime.kind === "managed_venv" && dataDirectory && isManagedPythonBackend(recipe.backend) ) { return dirname(managedVenvPython({ data_dir: dataDirectory }, recipe.backend)); } if ( (recipe.runtime.kind === "system" || recipe.runtime.kind === "binary") && recipe.runtime.ref.includes("/") ) { return dirname(recipe.runtime.ref); } return null;} export const pidExists = (pid: number): boolean => { try { process.kill(pid, 0); return true; } catch (error) { return (error as NodeJS.ErrnoException).code === "EPERM"; }}; export const buildProcessTree = (): Map<number, number[]> => { const result = spawnSync("ps", ["-eo", "pid=,ppid="]); if (result.status !== 0) { return new Map(); } const output = result.stdout.toString("utf-8").trim(); const tree = new Map<number, number[]>(); if (!output) { return tree; } for (const line of output.split("\n")) { const trimmed = line.trim(); const match = trimmed.match(/^(\d+)\s+(\d+)$/); if (!match) { continue; } const pid = Number(match[1]); const parent = Number(match[2]); const children = tree.get(parent) ?? []; children.push(pid); tree.set(parent, children); } return tree;}; export const collectChildren = ( tree: Map<number, number[]>, pid: number, accumulator: Set<number>,): void => { const children = tree.get(pid) ?? []; for (const child of children) { if (!accumulator.has(child)) { accumulator.add(child); collectChildren(tree, child, accumulator); } }};
import { spawnSync } from "node:child_process";
import { dirname } from "node:path";
import type { Recipe } from "../../models/types";
import type { Backend } from "@local-studio/contracts/recipes";
import { detectEngineFromArguments } from "../engine-spec";
import { extractFlag as extractFlagUtility } from "../argument-utilities";
import { isManagedPythonBackend, managedVenvPython } from "../runtimes/managed-venv";
import type { Config } from "../../../config/env";
export { extractFlagUtility as extractFlag };
const splitCommand = (command: string): string[] => {
const matches = command.match(/(?:[^\s"]+|"[^"]*")+/g) ?? [];
return matches.map((token) => token.replace(/^"|"$/g, ""));
};
export const detectBackend = (args: string[]): Backend | null => {
if (args.length === 0) return null;
return detectEngineFromArguments(args);
};
export const listProcesses = (): Array<{ pid: number; args: string[] }> => {
try {
const result = spawnSync("ps", ["-eo", "pid=,args="]);
if (result.status !== 0) {
return [];
}
const output = result.stdout.toString("utf-8").trim();
if (!output) {
return [];
}
return output
.split("\n")
.map((line) => {
const trimmed = line.trim();
const match = trimmed.match(/^(\d+)\s+(.*)$/);
if (!match) {
return { pid: 0, args: [] };
}
const pid = Number(match[1]);
const args = splitCommand(match[2] ?? "");
return { pid, args };
})
.filter((entry) => entry.pid > 0 && entry.args.length > 0);
} catch {
return [];
}
};
export const buildEnvironment = (
recipe: Recipe,
config?: Pick<Config, "data_dir">,
): Record<string, string> => {
const env: Record<string, string> = { ...process.env } as Record<string, string>;
env["FLASHINFER_DISABLE_VERSION_CHECK"] = "1";
const venvBin = resolveVenvBinForRecipe(recipe, config?.data_dir);
if (venvBin) {
env["PATH"] = `${venvBin}:${env["PATH"] ?? ""}`;
}
const environmentVariables: Record<string, string> = {};
if (recipe.env_vars && typeof recipe.env_vars === "object") {
for (const [key, value] of Object.entries(recipe.env_vars)) {
if (value !== undefined && value !== null) {
environmentVariables[String(key)] = String(value);
}
}
}
const extraEnvironment =
recipe.extra_args["env_vars"] || recipe.extra_args["env-vars"] || recipe.extra_args["envVars"];
if (extraEnvironment && typeof extraEnvironment === "object") {
for (const [key, value] of Object.entries(extraEnvironment as Record<string, unknown>)) {
if (value !== undefined && value !== null) {
environmentVariables[String(key)] = String(value);
}
}
}
for (const [key, value] of Object.entries(environmentVariables)) {
env[key] = value;
}
const readExtraArgument = (key: string): unknown => {
if (Object.prototype.hasOwnProperty.call(recipe.extra_args, key)) {
return recipe.extra_args[key];
}
const kebab = key.replace(/_/g, "-");
if (Object.prototype.hasOwnProperty.call(recipe.extra_args, kebab)) {
return recipe.extra_args[kebab];
}
const snake = key.replace(/-/g, "_");
if (Object.prototype.hasOwnProperty.call(recipe.extra_args, snake)) {
return recipe.extra_args[snake];
}
return undefined;
};
const isDefined = (value: unknown): boolean => {
return value !== undefined && value !== null && value !== false;
};
const visibleDevices =
readExtraArgument("visible_devices") ??
readExtraArgument("VISIBLE_DEVICES") ??
readExtraArgument("CUDA_VISIBLE_DEVICES") ??
readExtraArgument("cuda_visible_devices") ??
readExtraArgument("cuda-visible-devices");
const hipVisibleDevices =
readExtraArgument("hip_visible_devices") ?? readExtraArgument("HIP_VISIBLE_DEVICES");
const rocrVisibleDevices =
readExtraArgument("rocr_visible_devices") ?? readExtraArgument("ROCR_VISIBLE_DEVICES");
const forcedTool = (process.env["LOCAL_STUDIO_GPU_SMI_TOOL"] ?? "").trim().toLowerCase();
const platform =
forcedTool === "nvidia-smi"
? "cuda"
: forcedTool === "amd-smi" || forcedTool === "rocm-smi"
? "rocm"
: "unknown";
if (isDefined(visibleDevices)) {
const value = String(visibleDevices);
if (platform === "cuda") {
env["CUDA_VISIBLE_DEVICES"] = value;
} else if (platform === "rocm") {
env["HIP_VISIBLE_DEVICES"] = value;
env["ROCR_VISIBLE_DEVICES"] = value;
} else {
env["CUDA_VISIBLE_DEVICES"] = value;
env["HIP_VISIBLE_DEVICES"] = value;
env["ROCR_VISIBLE_DEVICES"] = value;
}
}
if (isDefined(hipVisibleDevices)) {
env["HIP_VISIBLE_DEVICES"] = String(hipVisibleDevices);
}
if (isDefined(rocrVisibleDevices)) {
env["ROCR_VISIBLE_DEVICES"] = String(rocrVisibleDevices);
}
return env;
};
function resolveVenvBinForRecipe(recipe: Recipe, dataDirectory?: string): string | null {
if (
recipe.runtime.kind === "managed_venv" &&
dataDirectory &&
isManagedPythonBackend(recipe.backend)
) {
return dirname(managedVenvPython({ data_dir: dataDirectory }, recipe.backend));
}
if (
(recipe.runtime.kind === "system" || recipe.runtime.kind === "binary") &&
recipe.runtime.ref.includes("/")
) {
return dirname(recipe.runtime.ref);
}
return null;
}
export const pidExists = (pid: number): boolean => {
try {
process.kill(pid, 0);
return true;
} catch (error) {
return (error as NodeJS.ErrnoException).code === "EPERM";
}
};
export const buildProcessTree = (): Map<number, number[]> => {
const result = spawnSync("ps", ["-eo", "pid=,ppid="]);
if (result.status !== 0) {
return new Map();
}
const output = result.stdout.toString("utf-8").trim();
const tree = new Map<number, number[]>();
if (!output) {
return tree;
}
for (const line of output.split("\n")) {
const trimmed = line.trim();
const match = trimmed.match(/^(\d+)\s+(\d+)$/);
if (!match) {
continue;
}
const pid = Number(match[1]);
const parent = Number(match[2]);
const children = tree.get(parent) ?? [];
children.push(pid);
tree.set(parent, children);
}
return tree;
};
export const collectChildren = (
tree: Map<number, number[]>,
pid: number,
accumulator: Set<number>,
): void => {
const children = tree.get(pid) ?? [];
for (const child of children) {
if (!accumulator.has(child)) {
accumulator.add(child);
collectChildren(tree, child, accumulator);
}
}
};