319 lines
10 KiB
TypeScript
319 lines
10 KiB
TypeScript
|
/**
|
||
|
|
* NVIDIA Driver Management
|
||
|
|
*
|
||
|
|
* Handles NVIDIA driver detection, installation, and container toolkit setup.
|
||
|
|
*/
|
||
|
|
|
||
|
|
import type { IDriverStatus } from '../interfaces/gpu.ts';
|
||
|
|
import { logger } from '../logger.ts';
|
||
|
|
import { BaseDriver, type IDriverInstallOptions } from './base-driver.ts';
|
||
|
|
|
||
|
|
/**
|
||
|
|
* NVIDIA Driver Manager
|
||
|
|
*/
|
||
|
|
export class NvidiaDriver extends BaseDriver {
|
||
|
|
public readonly vendor = 'nvidia' as const;
|
||
|
|
public readonly displayName = 'NVIDIA';
|
||
|
|
|
||
|
|
/**
|
||
|
|
* Check if the NVIDIA driver is installed
|
||
|
|
*/
|
||
|
|
public async isInstalled(): Promise<boolean> {
|
||
|
|
try {
|
||
|
|
const { stdout } = await this.execCommand('nvidia-smi --query-gpu=driver_version --format=csv,noheader', {
|
||
|
|
timeout: 5000,
|
||
|
|
ignoreErrors: true,
|
||
|
|
});
|
||
|
|
return stdout.trim().length > 0;
|
||
|
|
} catch {
|
||
|
|
return false;
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
/**
|
||
|
|
* Get NVIDIA driver status
|
||
|
|
*/
|
||
|
|
public async getStatus(): Promise<IDriverStatus> {
|
||
|
|
const status: IDriverStatus = {
|
||
|
|
vendor: 'nvidia',
|
||
|
|
installed: false,
|
||
|
|
containerSupport: false,
|
||
|
|
issues: [],
|
||
|
|
};
|
||
|
|
|
||
|
|
// Check if nvidia-smi is available
|
||
|
|
try {
|
||
|
|
const { stdout: driverVersion } = await this.execCommand(
|
||
|
|
'nvidia-smi --query-gpu=driver_version --format=csv,noheader | head -1',
|
||
|
|
{ timeout: 5000 },
|
||
|
|
);
|
||
|
|
status.installed = true;
|
||
|
|
status.version = driverVersion.trim();
|
||
|
|
} catch {
|
||
|
|
status.issues.push('NVIDIA driver not installed or nvidia-smi not available');
|
||
|
|
return status;
|
||
|
|
}
|
||
|
|
|
||
|
|
// Check CUDA toolkit
|
||
|
|
try {
|
||
|
|
const { stdout: cudaVersion } = await this.execCommand(
|
||
|
|
'nvcc --version 2>/dev/null | grep "release" | sed "s/.*release \\([0-9.]*\\).*/\\1/"',
|
||
|
|
{ timeout: 5000, ignoreErrors: true },
|
||
|
|
);
|
||
|
|
if (cudaVersion.trim()) {
|
||
|
|
status.toolkitVersion = cudaVersion.trim();
|
||
|
|
}
|
||
|
|
} catch {
|
||
|
|
// CUDA toolkit not installed
|
||
|
|
}
|
||
|
|
|
||
|
|
// Check nvidia-container-toolkit
|
||
|
|
try {
|
||
|
|
const { stdout: containerVersion } = await this.execCommand(
|
||
|
|
'nvidia-container-cli --version 2>&1 | head -1',
|
||
|
|
{ timeout: 5000, ignoreErrors: true },
|
||
|
|
);
|
||
|
|
if (containerVersion.includes('version')) {
|
||
|
|
status.containerSupport = true;
|
||
|
|
const match = containerVersion.match(/version (\d+\.\d+\.\d+)/);
|
||
|
|
if (match) {
|
||
|
|
status.containerRuntimeVersion = match[1];
|
||
|
|
}
|
||
|
|
}
|
||
|
|
} catch {
|
||
|
|
status.issues.push('NVIDIA Container Toolkit not installed');
|
||
|
|
}
|
||
|
|
|
||
|
|
// Check if Docker has nvidia runtime
|
||
|
|
try {
|
||
|
|
const { stdout: dockerInfo } = await this.execCommand(
|
||
|
|
'docker info --format "{{.Runtimes}}" 2>/dev/null',
|
||
|
|
{ timeout: 5000, ignoreErrors: true },
|
||
|
|
);
|
||
|
|
if (!dockerInfo.includes('nvidia')) {
|
||
|
|
status.issues.push('Docker nvidia runtime not configured');
|
||
|
|
}
|
||
|
|
} catch {
|
||
|
|
// Docker check failed
|
||
|
|
}
|
||
|
|
|
||
|
|
return status;
|
||
|
|
}
|
||
|
|
|
||
|
|
/**
|
||
|
|
* Install NVIDIA driver and optionally CUDA toolkit
|
||
|
|
*/
|
||
|
|
public async install(options: IDriverInstallOptions): Promise<boolean> {
|
||
|
|
if (!await this.isRoot()) {
|
||
|
|
logger.error('Root privileges required to install NVIDIA drivers');
|
||
|
|
return false;
|
||
|
|
}
|
||
|
|
|
||
|
|
const distro = await this.getLinuxDistro();
|
||
|
|
logger.info(`Detected Linux distribution: ${distro.id} ${distro.version}`);
|
||
|
|
|
||
|
|
try {
|
||
|
|
if (distro.id === 'ubuntu' || distro.id === 'debian') {
|
||
|
|
return await this.installOnDebian(options);
|
||
|
|
} else if (distro.id === 'fedora' || distro.id === 'rhel' || distro.id === 'centos' || distro.id === 'rocky' || distro.id === 'almalinux') {
|
||
|
|
return await this.installOnRhel(options);
|
||
|
|
} else {
|
||
|
|
logger.error(`Unsupported distribution: ${distro.id}`);
|
||
|
|
logger.info('Please install NVIDIA drivers manually');
|
||
|
|
return false;
|
||
|
|
}
|
||
|
|
} catch (error) {
|
||
|
|
logger.error(`Failed to install NVIDIA drivers: ${error instanceof Error ? error.message : String(error)}`);
|
||
|
|
return false;
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
/**
|
||
|
|
* Install on Debian/Ubuntu
|
||
|
|
*/
|
||
|
|
private async installOnDebian(options: IDriverInstallOptions): Promise<boolean> {
|
||
|
|
logger.info('Installing NVIDIA drivers on Debian/Ubuntu...');
|
||
|
|
|
||
|
|
// Add NVIDIA repository
|
||
|
|
await this.aptUpdate();
|
||
|
|
|
||
|
|
// Install prerequisites
|
||
|
|
await this.aptInstall(['software-properties-common', 'build-essential', 'dkms']);
|
||
|
|
|
||
|
|
// Add NVIDIA PPA (for Ubuntu)
|
||
|
|
try {
|
||
|
|
await this.execCommand('add-apt-repository -y ppa:graphics-drivers/ppa 2>/dev/null || true');
|
||
|
|
await this.aptUpdate();
|
||
|
|
} catch {
|
||
|
|
// PPA might not be available on all systems
|
||
|
|
}
|
||
|
|
|
||
|
|
// Install NVIDIA driver
|
||
|
|
const driverPackage = options.driverVersion
|
||
|
|
? `nvidia-driver-${options.driverVersion}`
|
||
|
|
: 'nvidia-driver-535'; // Default to stable version
|
||
|
|
|
||
|
|
await this.aptInstall(driverPackage);
|
||
|
|
|
||
|
|
// Install CUDA toolkit if requested
|
||
|
|
if (options.installToolkit) {
|
||
|
|
await this.installCudaToolkit(options);
|
||
|
|
}
|
||
|
|
|
||
|
|
// Install container support if requested
|
||
|
|
if (options.installContainerSupport) {
|
||
|
|
await this.installContainerSupport();
|
||
|
|
}
|
||
|
|
|
||
|
|
logger.success('NVIDIA driver installation completed');
|
||
|
|
logger.warn('A system reboot is required to load the new driver');
|
||
|
|
return true;
|
||
|
|
}
|
||
|
|
|
||
|
|
/**
|
||
|
|
* Install on RHEL/Fedora
|
||
|
|
*/
|
||
|
|
private async installOnRhel(options: IDriverInstallOptions): Promise<boolean> {
|
||
|
|
logger.info('Installing NVIDIA drivers on RHEL/Fedora...');
|
||
|
|
|
||
|
|
// Install prerequisites
|
||
|
|
await this.dnfInstall(['kernel-devel', 'kernel-headers', 'gcc', 'make', 'dkms', 'acpid']);
|
||
|
|
|
||
|
|
// Add NVIDIA CUDA repository
|
||
|
|
const distro = await this.getLinuxDistro();
|
||
|
|
const repoUrl = `https://developer.download.nvidia.com/compute/cuda/repos/rhel${distro.version.split('.')[0]}/x86_64/cuda-rhel${distro.version.split('.')[0]}.repo`;
|
||
|
|
|
||
|
|
await this.execCommand(`dnf config-manager --add-repo ${repoUrl}`);
|
||
|
|
|
||
|
|
// Install NVIDIA driver
|
||
|
|
await this.dnfInstall('nvidia-driver-latest-dkms');
|
||
|
|
|
||
|
|
// Install CUDA toolkit if requested
|
||
|
|
if (options.installToolkit) {
|
||
|
|
await this.dnfInstall('cuda');
|
||
|
|
}
|
||
|
|
|
||
|
|
// Install container support if requested
|
||
|
|
if (options.installContainerSupport) {
|
||
|
|
await this.installContainerSupport();
|
||
|
|
}
|
||
|
|
|
||
|
|
logger.success('NVIDIA driver installation completed');
|
||
|
|
logger.warn('A system reboot is required to load the new driver');
|
||
|
|
return true;
|
||
|
|
}
|
||
|
|
|
||
|
|
/**
|
||
|
|
* Install CUDA toolkit
|
||
|
|
*/
|
||
|
|
private async installCudaToolkit(options: IDriverInstallOptions): Promise<void> {
|
||
|
|
logger.info('Installing CUDA toolkit...');
|
||
|
|
|
||
|
|
const distro = await this.getLinuxDistro();
|
||
|
|
|
||
|
|
if (distro.id === 'ubuntu' || distro.id === 'debian') {
|
||
|
|
// Add CUDA repository
|
||
|
|
const cudaKeyUrl = 'https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2204/x86_64/cuda-keyring_1.1-1_all.deb';
|
||
|
|
await this.execCommand(`wget -q ${cudaKeyUrl} -O /tmp/cuda-keyring.deb && dpkg -i /tmp/cuda-keyring.deb`);
|
||
|
|
await this.aptUpdate();
|
||
|
|
|
||
|
|
const cudaPackage = options.toolkitVersion
|
||
|
|
? `cuda-toolkit-${options.toolkitVersion.replace('.', '-')}`
|
||
|
|
: 'cuda-toolkit';
|
||
|
|
|
||
|
|
await this.aptInstall(cudaPackage);
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
/**
|
||
|
|
* Install NVIDIA Container Toolkit
|
||
|
|
*/
|
||
|
|
public async installContainerSupport(): Promise<boolean> {
|
||
|
|
if (!await this.isRoot()) {
|
||
|
|
logger.error('Root privileges required to install NVIDIA Container Toolkit');
|
||
|
|
return false;
|
||
|
|
}
|
||
|
|
|
||
|
|
const distro = await this.getLinuxDistro();
|
||
|
|
logger.info('Installing NVIDIA Container Toolkit...');
|
||
|
|
|
||
|
|
try {
|
||
|
|
if (distro.id === 'ubuntu' || distro.id === 'debian') {
|
||
|
|
// Add repository
|
||
|
|
await this.execCommand(
|
||
|
|
'curl -fsSL https://nvidia.github.io/libnvidia-container/gpgkey | gpg --dearmor -o /usr/share/keyrings/nvidia-container-toolkit-keyring.gpg',
|
||
|
|
);
|
||
|
|
|
||
|
|
const distribution = `${distro.id}${distro.version}`;
|
||
|
|
await this.execCommand(
|
||
|
|
`curl -s -L https://nvidia.github.io/libnvidia-container/${distribution}/libnvidia-container.list | ` +
|
||
|
|
'sed "s#deb https://#deb [signed-by=/usr/share/keyrings/nvidia-container-toolkit-keyring.gpg] https://#g" | ' +
|
||
|
|
'tee /etc/apt/sources.list.d/nvidia-container-toolkit.list',
|
||
|
|
);
|
||
|
|
|
||
|
|
await this.aptUpdate();
|
||
|
|
await this.aptInstall('nvidia-container-toolkit');
|
||
|
|
} else {
|
||
|
|
// RHEL/Fedora
|
||
|
|
await this.execCommand(
|
||
|
|
'curl -s -L https://nvidia.github.io/libnvidia-container/stable/rpm/nvidia-container-toolkit.repo | ' +
|
||
|
|
'tee /etc/yum.repos.d/nvidia-container-toolkit.repo',
|
||
|
|
);
|
||
|
|
await this.dnfInstall('nvidia-container-toolkit');
|
||
|
|
}
|
||
|
|
|
||
|
|
// Configure Docker runtime
|
||
|
|
await this.configureDockerRuntime();
|
||
|
|
|
||
|
|
logger.success('NVIDIA Container Toolkit installed successfully');
|
||
|
|
return true;
|
||
|
|
} catch (error) {
|
||
|
|
logger.error(`Failed to install NVIDIA Container Toolkit: ${error instanceof Error ? error.message : String(error)}`);
|
||
|
|
return false;
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
/**
|
||
|
|
* Configure Docker to use NVIDIA runtime
|
||
|
|
*/
|
||
|
|
private async configureDockerRuntime(): Promise<void> {
|
||
|
|
logger.info('Configuring Docker to use NVIDIA runtime...');
|
||
|
|
|
||
|
|
try {
|
||
|
|
// Run nvidia-ctk to configure Docker
|
||
|
|
await this.execCommand('nvidia-ctk runtime configure --runtime=docker');
|
||
|
|
|
||
|
|
// Restart Docker
|
||
|
|
await this.execCommand('systemctl restart docker');
|
||
|
|
|
||
|
|
logger.success('Docker configured to use NVIDIA runtime');
|
||
|
|
} catch (error) {
|
||
|
|
logger.warn(`Could not configure Docker runtime automatically: ${error instanceof Error ? error.message : String(error)}`);
|
||
|
|
logger.info('Please run: nvidia-ctk runtime configure --runtime=docker');
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
/**
|
||
|
|
* Get available driver versions
|
||
|
|
*/
|
||
|
|
public async getAvailableVersions(): Promise<string[]> {
|
||
|
|
const versions: string[] = [];
|
||
|
|
|
||
|
|
try {
|
||
|
|
const distro = await this.getLinuxDistro();
|
||
|
|
|
||
|
|
if (distro.id === 'ubuntu' || distro.id === 'debian') {
|
||
|
|
const { stdout } = await this.execCommand(
|
||
|
|
'apt-cache search nvidia-driver | grep "^nvidia-driver-[0-9]" | sed "s/nvidia-driver-\\([0-9]*\\).*/\\1/" | sort -rn | uniq',
|
||
|
|
{ ignoreErrors: true },
|
||
|
|
);
|
||
|
|
versions.push(...stdout.trim().split('\n').filter((v: string) => v.trim()));
|
||
|
|
}
|
||
|
|
} catch {
|
||
|
|
// Failed to get versions
|
||
|
|
}
|
||
|
|
|
||
|
|
return versions;
|
||
|
|
}
|
||
|
|
}
|