chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,358 @@
|
||||
import fs from 'node:fs'
|
||||
import path from 'node:path'
|
||||
|
||||
import {
|
||||
NVIDIA_LIBS_PATH,
|
||||
NVIDIA_CUBLAS_PATH,
|
||||
NVIDIA_CUDNN_PATH,
|
||||
NVIDIA_CUDA_CUDART_PATH,
|
||||
NVIDIA_CUDA_CUPTI_PATH,
|
||||
NVIDIA_CUSPARSE_PATH,
|
||||
NVIDIA_CUSPARSELT_PATH,
|
||||
NVIDIA_CUSPARSE_FULL_PATH,
|
||||
NVIDIA_NCCL_PATH,
|
||||
NVIDIA_NVSHMEM_PATH,
|
||||
NVIDIA_NVJITLINK_PATH,
|
||||
NVIDIA_CUBLAS_MANIFEST_PATH,
|
||||
NVIDIA_CUDNN_MANIFEST_PATH,
|
||||
NVIDIA_CUDA_CUDART_MANIFEST_PATH,
|
||||
NVIDIA_CUDA_CUPTI_MANIFEST_PATH,
|
||||
NVIDIA_CUSPARSE_MANIFEST_PATH,
|
||||
NVIDIA_CUSPARSE_FULL_MANIFEST_PATH,
|
||||
NVIDIA_NCCL_MANIFEST_PATH,
|
||||
NVIDIA_NVSHMEM_MANIFEST_PATH,
|
||||
NVIDIA_NVJITLINK_MANIFEST_PATH,
|
||||
NVIDIA_CUDA_VERSION,
|
||||
NVIDIA_CUBLAS_VERSION,
|
||||
NVIDIA_CUDNN_VERSION,
|
||||
NVIDIA_CUDA_CUDART_VERSION,
|
||||
NVIDIA_CUDA_CUPTI_VERSION,
|
||||
NVIDIA_CUSPARSE_VERSION,
|
||||
NVIDIA_CUSPARSE_FULL_VERSION,
|
||||
NVIDIA_NCCL_VERSION,
|
||||
NVIDIA_NVSHMEM_VERSION,
|
||||
NVIDIA_NVJITLINK_VERSION,
|
||||
PYTORCH_NVIDIA_PATH,
|
||||
PYTORCH_TORCH_PATH
|
||||
} from '@/constants'
|
||||
import { FileHelper } from '@/helpers/file-helper'
|
||||
import { SystemHelper } from '@/helpers/system-helper'
|
||||
|
||||
import { createSetupStatus } from './setup-status'
|
||||
|
||||
const { type: OS_TYPE, cpuArchitecture: CPU_ARCH } =
|
||||
SystemHelper.getInformation()
|
||||
const NVIDIA_LIBRARY_LABELS = {
|
||||
cublas: 'cuBLAS',
|
||||
cudnn: 'cuDNN',
|
||||
cuda_cudart: 'CUDA Runtime',
|
||||
cuda_cupti: 'CUDA CUPTI',
|
||||
cusparse: 'cuSPARSE Lt',
|
||||
cusparse_full: 'cuSPARSE',
|
||||
nccl: 'NCCL',
|
||||
nvshmem: 'NVSHMEM',
|
||||
nvjitlink: 'nvJitLink'
|
||||
}
|
||||
|
||||
/**
|
||||
* Map CPU architecture to NVIDIA's architecture naming convention
|
||||
*/
|
||||
function mapToNvidiaArch(cpuArch) {
|
||||
// Map Node.js process.arch values to NVIDIA naming
|
||||
if (cpuArch === 'arm64' || cpuArch === 'aarch64') {
|
||||
return 'aarch64'
|
||||
}
|
||||
if (cpuArch === 'x64' || cpuArch === 'x86_64') {
|
||||
return 'x86_64'
|
||||
}
|
||||
|
||||
return 'x86_64'
|
||||
}
|
||||
|
||||
async function ensureDirectoryLink(linkPath, targetPath) {
|
||||
if (!fs.existsSync(targetPath)) {
|
||||
return
|
||||
}
|
||||
|
||||
await fs.promises.rm(linkPath, { recursive: true, force: true })
|
||||
await fs.promises.mkdir(path.dirname(linkPath), { recursive: true })
|
||||
|
||||
const relativeTarget = path.relative(path.dirname(linkPath), targetPath)
|
||||
const linkType = SystemHelper.isWindows() ? 'junction' : 'dir'
|
||||
|
||||
await fs.promises.symlink(relativeTarget, linkPath, linkType)
|
||||
}
|
||||
|
||||
async function ensureCompatibilityLinks() {
|
||||
await ensureDirectoryLink(NVIDIA_CUSPARSELT_PATH, NVIDIA_CUSPARSE_PATH)
|
||||
await ensureDirectoryLink(
|
||||
path.join(NVIDIA_LIBS_PATH, 'cuda_runtime'),
|
||||
NVIDIA_CUDA_CUDART_PATH
|
||||
)
|
||||
|
||||
if (fs.existsSync(PYTORCH_TORCH_PATH)) {
|
||||
await ensureDirectoryLink(PYTORCH_NVIDIA_PATH, NVIDIA_LIBS_PATH)
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Read manifest file to get installed version
|
||||
*/
|
||||
function readManifest(manifestPath) {
|
||||
if (!fs.existsSync(manifestPath)) {
|
||||
return null
|
||||
}
|
||||
|
||||
try {
|
||||
const content = fs.readFileSync(manifestPath, 'utf-8')
|
||||
|
||||
return JSON.parse(content)
|
||||
} catch {
|
||||
return null
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Get download URL for NVIDIA libraries
|
||||
*/
|
||||
function getNVIDIADownloadURL(library, version) {
|
||||
const ext = SystemHelper.isWindows() ? 'zip' : 'tar.xz'
|
||||
const arch = mapToNvidiaArch(CPU_ARCH)
|
||||
|
||||
// NVIDIA CDN URLs for CUDA libraries and more
|
||||
if (library === 'cublas') {
|
||||
return `https://developer.download.nvidia.com/compute/cuda/redist/libcublas/${OS_TYPE}-${arch}/libcublas-${OS_TYPE}-${arch}-${version}-archive.${ext}`
|
||||
} else if (library === 'cudnn') {
|
||||
return `https://developer.download.nvidia.com/compute/cudnn/redist/cudnn/${OS_TYPE}-${arch}/cudnn-${OS_TYPE}-${arch}-${version}_cuda${NVIDIA_CUDA_VERSION}-archive.${ext}`
|
||||
} else if (library === 'cuda_cudart') {
|
||||
return `https://developer.download.nvidia.com/compute/cuda/redist/cuda_cudart/${OS_TYPE}-${arch}/cuda_cudart-${OS_TYPE}-${arch}-${version}-archive.${ext}`
|
||||
} else if (library === 'cuda_cupti') {
|
||||
return `https://developer.download.nvidia.com/compute/cuda/redist/cuda_cupti/${OS_TYPE}-${arch}/cuda_cupti-${OS_TYPE}-${arch}-${version}-archive.${ext}`
|
||||
} else if (library === 'cusparse') {
|
||||
return `https://developer.download.nvidia.com/compute/cusparselt/redist/libcusparse_lt/${OS_TYPE}-${arch}/libcusparse_lt-${OS_TYPE}-${arch}-${version}_cuda${NVIDIA_CUDA_VERSION}-archive.${ext}`
|
||||
} else if (library === 'cusparse_full') {
|
||||
return `https://developer.download.nvidia.com/compute/cuda/redist/libcusparse/${OS_TYPE}-${arch}/libcusparse-${OS_TYPE}-${arch}-${version}-archive.${ext}`
|
||||
} else if (library === 'nccl') {
|
||||
// NCCL is only available on Linux x86_64
|
||||
if (!SystemHelper.isLinux() || arch !== 'x86_64') {
|
||||
throw new Error('NCCL is only available on Linux x86_64')
|
||||
}
|
||||
|
||||
return `https://developer.download.nvidia.com/compute/nccl/redist/nccl/${OS_TYPE}-${arch}/nccl-${OS_TYPE}-${arch}-${version}-archive.${ext}`
|
||||
} else if (library === 'nvshmem') {
|
||||
// NVSHMEM is only available on Linux x86_64
|
||||
if (!SystemHelper.isLinux() || arch !== 'x86_64') {
|
||||
throw new Error('NVSHMEM is only available on Linux x86_64')
|
||||
}
|
||||
|
||||
return `https://developer.download.nvidia.com/compute/nvshmem/redist/libnvshmem/${OS_TYPE}-${arch}/libnvshmem-${OS_TYPE}-${arch}-${version}_cuda${NVIDIA_CUDA_VERSION}-archive.${ext}`
|
||||
} else if (library === 'nvjitlink') {
|
||||
return `https://developer.download.nvidia.com/compute/cuda/redist/libnvjitlink/${OS_TYPE}-${arch}/libnvjitlink-${OS_TYPE}-${arch}-${version}-archive.${ext}`
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Install NVIDIA libraries if needed
|
||||
*/
|
||||
async function installNVIDIALibrary(
|
||||
library,
|
||||
requiredVersion,
|
||||
targetPath,
|
||||
manifestPath
|
||||
) {
|
||||
const libraryLabel = NVIDIA_LIBRARY_LABELS[library] || library
|
||||
const status = createSetupStatus(`Setting up ${libraryLabel}...`).start()
|
||||
const manifest = readManifest(manifestPath)
|
||||
|
||||
if (!manifest || manifest.version !== requiredVersion) {
|
||||
const ext = SystemHelper.isWindows() ? 'zip' : 'tar.xz'
|
||||
const archivePath = path.join(
|
||||
NVIDIA_LIBS_PATH,
|
||||
`${library}-${requiredVersion}.${ext}`
|
||||
)
|
||||
|
||||
// Clean up old version
|
||||
await fs.promises.rm(targetPath, { recursive: true, force: true })
|
||||
await fs.promises.rm(archivePath, { force: true })
|
||||
|
||||
// Create target directory
|
||||
await fs.promises.mkdir(targetPath, { recursive: true })
|
||||
|
||||
try {
|
||||
const downloadURL = getNVIDIADownloadURL(library, requiredVersion)
|
||||
status.pause()
|
||||
|
||||
await FileHelper.downloadFile(downloadURL, archivePath, {
|
||||
cliProgress: true,
|
||||
parallelStreams: 3,
|
||||
skipExisting: false
|
||||
})
|
||||
status.text = `Installing ${libraryLabel}...`
|
||||
status.start()
|
||||
|
||||
// Extract archive using unified method
|
||||
await FileHelper.extractArchive(archivePath, targetPath, {
|
||||
stripComponents: 1
|
||||
})
|
||||
|
||||
// Clean up and create manifest
|
||||
await Promise.all([
|
||||
fs.promises.rm(archivePath, { force: true }),
|
||||
FileHelper.createManifestFile(manifestPath, library, requiredVersion, {
|
||||
os: SystemHelper.getInformation().type,
|
||||
architecture: SystemHelper.getInformation().cpuArchitecture
|
||||
})
|
||||
])
|
||||
status.succeed(`${libraryLabel} ${requiredVersion} ready`)
|
||||
} catch (error) {
|
||||
if (status.isSpinning) {
|
||||
status.fail(`Failed to set up ${libraryLabel}`)
|
||||
}
|
||||
throw new Error(
|
||||
`${libraryLabel} may require manual download from https://developer.nvidia.com/cuda-downloads: ${error}`
|
||||
)
|
||||
}
|
||||
} else {
|
||||
status.succeed(`${libraryLabel}: ${requiredVersion}`)
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Main setup function
|
||||
*/
|
||||
async function setupNVIDIALibs() {
|
||||
// Skip on macOS since there is no CUDA involved
|
||||
if (SystemHelper.isMacOS()) {
|
||||
return
|
||||
}
|
||||
const status = createSetupStatus('Checking CUDA runtime support...').start()
|
||||
|
||||
try {
|
||||
const { getLlama, LlamaLogLevel } = await Function(
|
||||
'return import("node-llama-cpp")'
|
||||
)()
|
||||
const llama = await getLlama({
|
||||
logLevel: LlamaLogLevel.disabled
|
||||
})
|
||||
|
||||
const hasGPU = await SystemHelper.hasGPU(llama)
|
||||
|
||||
if (!hasGPU) {
|
||||
status.succeed('CUDA runtime: skipped')
|
||||
return
|
||||
}
|
||||
|
||||
status.succeed('CUDA runtime: detected')
|
||||
|
||||
// Install/update cuBLAS
|
||||
await installNVIDIALibrary(
|
||||
'cublas',
|
||||
NVIDIA_CUBLAS_VERSION,
|
||||
NVIDIA_CUBLAS_PATH,
|
||||
NVIDIA_CUBLAS_MANIFEST_PATH
|
||||
)
|
||||
|
||||
// Install/update cuDNN
|
||||
await installNVIDIALibrary(
|
||||
'cudnn',
|
||||
NVIDIA_CUDNN_VERSION,
|
||||
NVIDIA_CUDNN_PATH,
|
||||
NVIDIA_CUDNN_MANIFEST_PATH
|
||||
)
|
||||
|
||||
// Install/update CUDA cudart runtime
|
||||
await installNVIDIALibrary(
|
||||
'cuda_cudart',
|
||||
NVIDIA_CUDA_CUDART_VERSION,
|
||||
NVIDIA_CUDA_CUDART_PATH,
|
||||
NVIDIA_CUDA_CUDART_MANIFEST_PATH
|
||||
)
|
||||
|
||||
// Install/update CUDA CUPTI
|
||||
await installNVIDIALibrary(
|
||||
'cuda_cupti',
|
||||
NVIDIA_CUDA_CUPTI_VERSION,
|
||||
NVIDIA_CUDA_CUPTI_PATH,
|
||||
NVIDIA_CUDA_CUPTI_MANIFEST_PATH
|
||||
)
|
||||
|
||||
// Install/update cuSPARSE-Lt (Linux only, both x86_64 and aarch64)
|
||||
if (SystemHelper.isLinux()) {
|
||||
try {
|
||||
await installNVIDIALibrary(
|
||||
'cusparse',
|
||||
NVIDIA_CUSPARSE_VERSION,
|
||||
NVIDIA_CUSPARSE_PATH,
|
||||
NVIDIA_CUSPARSE_MANIFEST_PATH
|
||||
)
|
||||
} catch (error) {
|
||||
status.warn(`cuSPARSE Lt skipped: ${error.message}`)
|
||||
}
|
||||
}
|
||||
|
||||
// Install/update cuSPARSE (Linux only, both x86_64 and aarch64)
|
||||
if (SystemHelper.isLinux()) {
|
||||
try {
|
||||
await installNVIDIALibrary(
|
||||
'cusparse_full',
|
||||
NVIDIA_CUSPARSE_FULL_VERSION,
|
||||
NVIDIA_CUSPARSE_FULL_PATH,
|
||||
NVIDIA_CUSPARSE_FULL_MANIFEST_PATH
|
||||
)
|
||||
} catch (error) {
|
||||
status.warn(`cuSPARSE skipped: ${error.message}`)
|
||||
}
|
||||
}
|
||||
|
||||
// Install/update nvJitLink (Linux only)
|
||||
if (SystemHelper.isLinux()) {
|
||||
try {
|
||||
await installNVIDIALibrary(
|
||||
'nvjitlink',
|
||||
NVIDIA_NVJITLINK_VERSION,
|
||||
NVIDIA_NVJITLINK_PATH,
|
||||
NVIDIA_NVJITLINK_MANIFEST_PATH
|
||||
)
|
||||
} catch (error) {
|
||||
status.warn(`nvJitLink skipped: ${error.message}`)
|
||||
}
|
||||
}
|
||||
|
||||
// Install/update NCCL (Linux x86_64 only)
|
||||
if (SystemHelper.isLinux() && mapToNvidiaArch(CPU_ARCH) === 'x86_64') {
|
||||
try {
|
||||
await installNVIDIALibrary(
|
||||
'nccl',
|
||||
NVIDIA_NCCL_VERSION,
|
||||
NVIDIA_NCCL_PATH,
|
||||
NVIDIA_NCCL_MANIFEST_PATH
|
||||
)
|
||||
} catch (error) {
|
||||
status.warn(`NCCL skipped: ${error.message}`)
|
||||
}
|
||||
}
|
||||
|
||||
// Install/update NVSHMEM (Linux x86_64 only)
|
||||
if (SystemHelper.isLinux() && mapToNvidiaArch(CPU_ARCH) === 'x86_64') {
|
||||
try {
|
||||
await installNVIDIALibrary(
|
||||
'nvshmem',
|
||||
NVIDIA_NVSHMEM_VERSION,
|
||||
NVIDIA_NVSHMEM_PATH,
|
||||
NVIDIA_NVSHMEM_MANIFEST_PATH
|
||||
)
|
||||
} catch (error) {
|
||||
status.warn(`NVSHMEM skipped: ${error.message}`)
|
||||
}
|
||||
}
|
||||
|
||||
await ensureCompatibilityLinks()
|
||||
} catch (error) {
|
||||
if (status.isSpinning) {
|
||||
status.fail('Failed to inspect CUDA runtime support')
|
||||
}
|
||||
throw error
|
||||
}
|
||||
}
|
||||
|
||||
export default setupNVIDIALibs
|
||||
Reference in New Issue
Block a user