362 lines
11 KiB
TypeScript
362 lines
11 KiB
TypeScript
import _ from 'lodash';
|
|
|
|
export const GPUDriverMap = {
|
|
NVIDIA: 'cuda',
|
|
AMD: 'rocm',
|
|
ASCEND: 'cann',
|
|
HYGON: 'dtk',
|
|
MOORE_THREADS: 'musa',
|
|
ILUVATAR: 'corex',
|
|
CAMBRICON: 'neuware',
|
|
METAX: 'maca',
|
|
THEAD: 'hggc'
|
|
};
|
|
|
|
export const ManufacturerMap = {
|
|
[GPUDriverMap.NVIDIA]: 'NVIDIA',
|
|
[GPUDriverMap.AMD]: 'AMD',
|
|
[GPUDriverMap.ASCEND]: 'vendor.ascend',
|
|
[GPUDriverMap.HYGON]: 'vendor.hygon',
|
|
[GPUDriverMap.MOORE_THREADS]: 'vendor.moorthreads',
|
|
[GPUDriverMap.ILUVATAR]: 'vendor.iluvatar',
|
|
[GPUDriverMap.CAMBRICON]: 'vendor.cambricon',
|
|
[GPUDriverMap.METAX]: 'vendor.metax',
|
|
[GPUDriverMap.THEAD]: 'vendor.thead'
|
|
};
|
|
|
|
export const GPUsConfigs: Record<
|
|
string,
|
|
{ label: string; value: string; runtime: string; driver: string }
|
|
> = {
|
|
[GPUDriverMap.NVIDIA]: {
|
|
label: 'Nvidia', // this label is used in echo command, do not translate
|
|
value: GPUDriverMap.NVIDIA,
|
|
runtime: 'nvidia',
|
|
driver: 'nvidia-smi'
|
|
},
|
|
[GPUDriverMap.AMD]: {
|
|
label: 'AMD',
|
|
value: GPUDriverMap.AMD,
|
|
runtime: 'amd',
|
|
driver: 'amd-smi static'
|
|
},
|
|
[GPUDriverMap.ASCEND]: {
|
|
label: 'Ascend',
|
|
value: GPUDriverMap.ASCEND,
|
|
runtime: 'ascend',
|
|
driver: 'npu-smi info'
|
|
},
|
|
[GPUDriverMap.HYGON]: {
|
|
label: 'Hygon',
|
|
value: GPUDriverMap.HYGON,
|
|
runtime: '', // TODO: confirm runtime name
|
|
driver: 'hy-smi'
|
|
},
|
|
[GPUDriverMap.MOORE_THREADS]: {
|
|
label: 'Moore Threads',
|
|
value: GPUDriverMap.MOORE_THREADS,
|
|
runtime: 'mthreads',
|
|
driver: 'mthreads-gmi'
|
|
},
|
|
[GPUDriverMap.ILUVATAR]: {
|
|
label: 'Iluvatar',
|
|
value: GPUDriverMap.ILUVATAR,
|
|
runtime: 'iluvatar',
|
|
driver: 'ixsmi'
|
|
},
|
|
[GPUDriverMap.CAMBRICON]: {
|
|
label: 'Cambricon',
|
|
value: GPUDriverMap.CAMBRICON,
|
|
runtime: 'cambricon',
|
|
driver: 'cnmon info'
|
|
},
|
|
[GPUDriverMap.METAX]: {
|
|
label: 'Metax',
|
|
value: GPUDriverMap.METAX,
|
|
runtime: 'metax',
|
|
driver: 'mx-smi'
|
|
},
|
|
[GPUDriverMap.THEAD]: {
|
|
label: 'T-Head PPU',
|
|
value: GPUDriverMap.THEAD,
|
|
runtime: '', // TODO: confirm runtime name
|
|
driver: 'ppu-smi'
|
|
}
|
|
};
|
|
|
|
// Generate the Docker environment command: available for NVIDIA, AMD, Ascend, Moore Threads
|
|
const generateNvidiaDockerEnvCommand = (config: {
|
|
label: string;
|
|
value: string;
|
|
runtime: string;
|
|
driver: string;
|
|
}) => {
|
|
const runtimeCheck =
|
|
config.runtime === 'nvidia' ? 'Runtime.*nvidia' : config.runtime;
|
|
|
|
return `${config.driver} >/dev/null 2>&1 && echo "${config.label} driver OK" || (echo "${config.label} driver issue"; exit 1) && sudo docker info 2>/dev/null | grep -q "${runtimeCheck}" && echo "${config.label} Container Toolkit OK" || (echo "${config.label} Container Toolkit not configured"; exit 1)`;
|
|
};
|
|
|
|
// avaliable for Hygon、Iluvatar、MetaX 、Cambricon、T-Head
|
|
const generateHygonDockerEnvCommand = (config: {
|
|
label: string;
|
|
value: string;
|
|
runtime: string;
|
|
driver: string;
|
|
}) => {
|
|
return `${config.driver} >/dev/null 2>&1 && echo "${config.label} driver OK" || (echo "${config.label} driver issue"; exit 1)`;
|
|
};
|
|
|
|
// available for Kubernetes
|
|
export const generateKubernetesEnvCommand = (gpu: string) => {
|
|
const config = GPUsConfigs[gpu] || {};
|
|
return `kubectl get runtimeclass ${config.runtime} > /dev/null 2>&1 && echo "${config.label} runtimeclass registered" || (echo "${config.label} runtimeclass issue"; exit 1)`;
|
|
};
|
|
|
|
export const dockerEnvCommandMap = {
|
|
[GPUDriverMap.NVIDIA]: generateNvidiaDockerEnvCommand(
|
|
GPUsConfigs[GPUDriverMap.NVIDIA]
|
|
),
|
|
[GPUDriverMap.AMD]: generateNvidiaDockerEnvCommand(
|
|
GPUsConfigs[GPUDriverMap.AMD]
|
|
),
|
|
[GPUDriverMap.ASCEND]: generateNvidiaDockerEnvCommand(
|
|
GPUsConfigs[GPUDriverMap.ASCEND]
|
|
),
|
|
[GPUDriverMap.MOORE_THREADS]: generateNvidiaDockerEnvCommand(
|
|
GPUsConfigs[GPUDriverMap.MOORE_THREADS]
|
|
),
|
|
[GPUDriverMap.HYGON]: generateHygonDockerEnvCommand(
|
|
GPUsConfigs[GPUDriverMap.HYGON]
|
|
),
|
|
[GPUDriverMap.ILUVATAR]: generateHygonDockerEnvCommand(
|
|
GPUsConfigs[GPUDriverMap.ILUVATAR]
|
|
),
|
|
[GPUDriverMap.CAMBRICON]: generateHygonDockerEnvCommand(
|
|
GPUsConfigs[GPUDriverMap.CAMBRICON]
|
|
),
|
|
[GPUDriverMap.METAX]: generateHygonDockerEnvCommand(
|
|
GPUsConfigs[GPUDriverMap.METAX]
|
|
),
|
|
[GPUDriverMap.THEAD]: generateHygonDockerEnvCommand(
|
|
GPUsConfigs[GPUDriverMap.THEAD]
|
|
)
|
|
};
|
|
|
|
interface AddWorkerCommandParams {
|
|
server: string;
|
|
tag: string;
|
|
token: string;
|
|
image: string;
|
|
gpu: string;
|
|
workerIP?: string;
|
|
modelDir?: string;
|
|
containerName?: string;
|
|
gpustackDataVolume?: string;
|
|
cacheDir?: string;
|
|
}
|
|
|
|
const generateEnvArgs = (params: any) => {
|
|
const registrationInfo = {
|
|
...params.registrationInfo,
|
|
env: {
|
|
...(params.registrationInfo?.env || {}),
|
|
...params.extraEnv
|
|
}
|
|
};
|
|
// generate environment variables args from registrationInfo.env
|
|
let envArgs = '';
|
|
if (
|
|
registrationInfo.env &&
|
|
typeof registrationInfo.env === 'object' &&
|
|
Object.keys(registrationInfo.env).length > 0
|
|
) {
|
|
Object.keys(registrationInfo.env).forEach((key) => {
|
|
const value = registrationInfo.env[key];
|
|
envArgs += `-e "${key}=${value}" \\\n `;
|
|
});
|
|
}
|
|
return envArgs;
|
|
};
|
|
|
|
// concat the args, the args is a key-value
|
|
const generateExtraArgs = (params: any) => {
|
|
const args = params.registrationInfo?.args || [];
|
|
if (args.length === 0) {
|
|
return '';
|
|
}
|
|
let argsStr = '';
|
|
args.forEach((item: string[]) => {
|
|
argsStr += `${item[0]} ${_.isBoolean(item[1]) ? item[1] : item[1] || ''} \\\n `;
|
|
});
|
|
return argsStr;
|
|
};
|
|
|
|
const generateExtraModelDirArg = (modelDir: string) => {
|
|
const pathList = modelDir
|
|
?.split(',')
|
|
.map((item) => item.trim())
|
|
.filter((item) => item);
|
|
if (pathList && pathList.length > 0) {
|
|
return pathList
|
|
.map((item) => `--volume ${item}:${item} \\`)
|
|
.join('\n ');
|
|
}
|
|
return '';
|
|
};
|
|
|
|
const setNormalArgs = (params: any) => {
|
|
return `sudo docker run -d --name ${params.containerName || 'gpustack-worker'} \\
|
|
-e "GPUSTACK_RUNTIME_DEPLOY_MIRRORED_NAME=${params.containerName || 'gpustack-worker'}" \\
|
|
${generateEnvArgs(params)}
|
|
--restart=unless-stopped \\
|
|
--privileged \\
|
|
--network=host \\
|
|
--volume /var/run/docker.sock:/var/run/docker.sock \\
|
|
--volume ${params.gpustackDataVolume || 'gpustack-data'}:/var/lib/gpustack \\
|
|
${generateExtraModelDirArg(params.modelDir)}
|
|
${params.cacheDir ? `--volume ${params.cacheDir}:/var/lib/gpustack/cache \\` : ''}`;
|
|
};
|
|
|
|
const setWorkerIPArg = (params: any) => {
|
|
return `${params.workerIP ? `--worker-ip ${params.workerIP} \\` : ''}
|
|
${params.advertisAddress ? `--advertise-address ${params.advertisAddress} \\` : ''}`;
|
|
};
|
|
|
|
const setImageArgs = (params: any) => {
|
|
return `${params.image} \\
|
|
--server-url ${params.server} \\`;
|
|
};
|
|
|
|
// avaliable for NVIDIA、MThreads
|
|
const registerWorker = (params: AddWorkerCommandParams) => {
|
|
const config = GPUsConfigs[params.gpu];
|
|
const commonArgs = setNormalArgs(params);
|
|
const imageArgs = setImageArgs(params);
|
|
|
|
// remove empty enter lines and trailing backslash
|
|
return `${commonArgs}
|
|
--runtime ${config.runtime} \\
|
|
${imageArgs}
|
|
${setWorkerIPArg(params)}`;
|
|
};
|
|
|
|
// avaliable for AMD
|
|
const registerAMDWorker = (params: AddWorkerCommandParams) => {
|
|
const config = GPUsConfigs[params.gpu];
|
|
const commonArgs = setNormalArgs(params);
|
|
const imageArgs = setImageArgs(params);
|
|
// remove empty enter lines and trailing backslash
|
|
return `${commonArgs}
|
|
--volume /opt/rocm:/opt/rocm:ro \\
|
|
--runtime ${config.runtime} \\
|
|
${imageArgs}
|
|
${setWorkerIPArg(params)}`;
|
|
};
|
|
|
|
// avaliable for Ascend
|
|
const registerAscendWorker = (params: AddWorkerCommandParams) => {
|
|
const config = GPUsConfigs[params.gpu];
|
|
const commonArgs = setNormalArgs(params);
|
|
const imageArgs = setImageArgs(params);
|
|
return `${commonArgs}
|
|
--env "ASCEND_VISIBLE_DEVICES=$(sudo ls /dev/davinci* | head -1 | grep -o '[0-9]\\+' || echo "0")" \\
|
|
--runtime ${config.runtime} \\
|
|
${imageArgs}
|
|
${setWorkerIPArg(params)}`;
|
|
};
|
|
|
|
// avaliable for t-head
|
|
|
|
const registerTHeadWorker = (params: AddWorkerCommandParams) => {
|
|
const config = GPUsConfigs[params.gpu];
|
|
const commonArgs = setNormalArgs({
|
|
...params,
|
|
extraEnv: {
|
|
GPUSTACK_RUNTIME_DOCKER_RESOURCE_INJECTION_POLICY: 'CDI'
|
|
}
|
|
});
|
|
const imageArgs = setImageArgs(params);
|
|
return `${commonArgs}
|
|
--volume /usr/local/PPU_SDK:/usr/local/PPU_SDK:ro \\
|
|
--volume /var/run/cdi:/var/run/cdi \\
|
|
${imageArgs}
|
|
${setWorkerIPArg(params)}`;
|
|
};
|
|
|
|
const registerHygonWorker = (params: AddWorkerCommandParams) => {
|
|
const config = GPUsConfigs[params.gpu];
|
|
const commonArgs = setNormalArgs(params);
|
|
const imageArgs = setImageArgs(params);
|
|
return `${commonArgs}
|
|
--volume /opt/hyhal:/opt/hyhal:ro \\
|
|
--volume /opt/dtk:/opt/dtk:ro \\
|
|
--env ROCM_PATH=/opt/dtk \\
|
|
--env ROCM_SMI_LIB_PATH=/opt/hyhal/lib \\
|
|
${imageArgs}
|
|
${setWorkerIPArg(params)}`;
|
|
};
|
|
|
|
const registerIluvatarWorker = (params: AddWorkerCommandParams) => {
|
|
const config = GPUsConfigs[params.gpu];
|
|
const commonArgs = setNormalArgs(params);
|
|
const imageArgs = setImageArgs(params);
|
|
return `${commonArgs}
|
|
--volume /usr/local/corex:/usr/local/corex:ro \\
|
|
--runtime ${config.runtime} \\
|
|
${imageArgs}
|
|
${setWorkerIPArg(params)}`;
|
|
};
|
|
|
|
const registerMetaXWorker = (params: AddWorkerCommandParams) => {
|
|
const config = GPUsConfigs[params.gpu];
|
|
const commonArgs = setNormalArgs(params);
|
|
const imageArgs = setImageArgs(params);
|
|
return `${commonArgs}
|
|
--volume /opt/mxdriver:/opt/mxdriver:ro \\
|
|
--volume /opt/maca:/opt/maca:ro \\
|
|
${imageArgs}
|
|
${setWorkerIPArg(params)}`;
|
|
};
|
|
|
|
const registerCambriconWorker = (params: AddWorkerCommandParams) => {
|
|
const config = GPUsConfigs[params.gpu];
|
|
const commonArgs = setNormalArgs(params);
|
|
const imageArgs = setImageArgs(params);
|
|
return `${commonArgs}
|
|
--volume /usr/local/neuware:/usr/local/neuware:ro \\
|
|
--volume /usr/bin/cnmon:/usr/bin/cnmon \\
|
|
${imageArgs}
|
|
${setWorkerIPArg(params)}`;
|
|
};
|
|
|
|
export const registerAddWokerCommandMap = {
|
|
[GPUDriverMap.NVIDIA]: registerWorker,
|
|
[GPUDriverMap.AMD]: registerAMDWorker,
|
|
[GPUDriverMap.ASCEND]: registerAscendWorker,
|
|
[GPUDriverMap.HYGON]: registerHygonWorker,
|
|
[GPUDriverMap.ILUVATAR]: registerIluvatarWorker,
|
|
[GPUDriverMap.CAMBRICON]: registerCambriconWorker,
|
|
[GPUDriverMap.METAX]: registerMetaXWorker,
|
|
[GPUDriverMap.MOORE_THREADS]: registerWorker,
|
|
[GPUDriverMap.THEAD]: registerTHeadWorker
|
|
};
|
|
|
|
export const AddWorkerDockerNotes: Record<string, string[]> = {
|
|
[GPUDriverMap.NVIDIA]: ['clusters.addworker.nvidiaNotes'],
|
|
[GPUDriverMap.AMD]: ['clusters.addworker.amdNotes-01'],
|
|
[GPUDriverMap.MOORE_THREADS]: [],
|
|
[GPUDriverMap.ASCEND]: [],
|
|
[GPUDriverMap.HYGON]: [
|
|
'clusters.addworker.hygonNotes',
|
|
'clusters.addworker.hygonNotes-02'
|
|
],
|
|
[GPUDriverMap.ILUVATAR]: ['clusters.addworker.corexNotes'],
|
|
[GPUDriverMap.CAMBRICON]: ['clusters.addworker.cambriconNotes'],
|
|
[GPUDriverMap.METAX]: ['clusters.addworker.metaxNotes'],
|
|
[GPUDriverMap.THEAD]: [
|
|
'clusters.addworker.theadNotes',
|
|
'clusters.addworker.theadNotes-02'
|
|
]
|
|
};
|