# Copyright 2021 Google LLC # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. import argparse import atexit import os import pathlib import re import shlex import subprocess import sys from datetime import datetime from enum import Enum, auto from typing import Optional DRIVER_VERSION = "525.125.06" K80_DRIVER_VERSION = "470.199.02" DRIVER_URL = f"https://us.download.nvidia.com/tesla/{DRIVER_VERSION}/NVIDIA-Linux-x86_64-{DRIVER_VERSION}.run" K80_DRIVER_URL = f"https://us.download.nvidia.com/tesla/{K80_DRIVER_VERSION}/NVIDIA-Linux-x86_64-{K80_DRIVER_VERSION}.run" TESLA_K80_DEVICE_CODE = "10de:102d" class System(Enum): CentOS = auto() Debian = auto() Fedora = auto() RHEL = auto() Rocky = auto() SUSE = auto() Ubuntu = auto() # CentOS 7 and RHEL 7 may require Python3 to be installed before this script can be run. SUPPORTED_SYSTEMS = { # CentOS 8 is dead: https://www.centos.org/centos-linux-eol/, but there's CentOS Stream 8 System.CentOS: {"7", "8"}, System.Debian: {"10", "11"}, System.Fedora: set(), System.RHEL: {"7", "8", "9"}, System.Rocky: {"8", "9"}, System.SUSE: set(), System.Ubuntu: {"18", "20", "21", "22"} } INSTALLER_DIR = pathlib.Path('/opt/google/gpu-installer/') DEPENDENCIES_INSTALLED_FLAG = INSTALLER_DIR / 'deps_installed.flag' INSTALLER_DIR.mkdir(parents=True, exist_ok=True) class Logger: STDOUT_LOG = INSTALLER_DIR / 'out.log' STDOUT_LOG_F = None STDERR_LOG = INSTALLER_DIR / 'err.log' STDERR_LOG_F = None @classmethod def close_logs(cls): if cls.STDOUT_LOG_F: cls.STDOUT_LOG_F.close() if cls.STDERR_LOG_F: cls.STDERR_LOG_F.close() @classmethod def setup_log_dir(cls): """ Create the LOG_DIR path and STD(OUT|ERR)_LOG files. """ cls.STDOUT_LOG.touch(exist_ok=True) cls.STDERR_LOG.touch(exist_ok=True) cls.STDOUT_LOG_F = open(cls.STDOUT_LOG, mode='a') cls.STDERR_LOG_F = open(cls.STDERR_LOG, mode='a') atexit.register(cls.close_logs) @classmethod def print_out(cls, msg: str, end=os.linesep, print_=True): if cls.STDOUT_LOG_F: cls.STDOUT_LOG_F.write(msg + end) cls.STDOUT_LOG_F.flush() if print_: print(msg, end=end, file=sys.stdout) @classmethod def print_err(cls, msg: str, end=os.linesep, print_=True): if cls.STDERR_LOG_F: cls.STDERR_LOG_F.write(msg + end) cls.STDERR_LOG_F.flush() if print_: print(msg, end=end, file=sys.stderr) print_out = Logger.print_out print_err = Logger.print_err def run(command: str, check=True, input=None, cwd=None, silent=False, environment=None, retries=0) -> subprocess.CompletedProcess: """ Runs a provided command, streaming its output to the log files. :param command: A command to be executed, as a single string. :param check: If true, will throw exception on failure (exit code != 0) :param input: Input for the executed command. :param cwd: Directory in which to execute the command. :param silent: If set to True, the output of command won't be logged or printed. :param environment: A set of environment variable for the process to use. If None, the current env is inherited. :param retries: How many times should the command be repeated if it exits with non-zero code. :return: CompletedProcess instance - the result of the command execution. """ if not silent: log_msg = f"[{datetime.now().strftime('%Y-%m-%d %H:%M:%S')}] " \ f"Executing: {command}" + os.linesep print_out(log_msg) print_err(log_msg, print_=False) try_count = 0 while try_count <= retries: try: proc = subprocess.run(shlex.split(command), check=check, stderr=subprocess.PIPE, stdout=subprocess.PIPE, input=input, cwd=cwd, env=environment) except subprocess.SubprocessError as err: print_err(f"Error while executing `{command}`:") print_err(str(err)) if try_count == retries: raise err else: break try_count += 1 if not silent: print_err(proc.stderr.decode()) print_out(proc.stdout.decode()) return proc def detect_gpu_device() -> Optional[str]: """ Check if there is an NVIDIA GPU device attached and return its device code. """ lspci = run('lspci -n') output = lspci.stdout.decode() dev_re = re.compile(r"10de:[\w\d]{4}") for line in output.splitlines(): dev_code = dev_re.findall(line) if dev_code: return dev_code[0] else: return None def check_python_version(): """ Makes sure that the script is run with Python 3.6 or newer. """ if sys.version_info.major == 3 and sys.version_info.minor >= 6: return version = "{}.{}".format(sys.version_info.major, sys.version_info.minor) raise RuntimeError("Unsupported Python version {}. " "Supported versions: 3.6 - 3.11".format(version)) def detect_linux_distro() -> (System, str): """ Checks the /etc/os-release file to figure out what distribution of OS we're running. """ with open('/etc/os-release') as os_release: lines = [line.strip() for line in os_release.readlines() if line.strip() != ''] info = {k: v.strip("'\"") for k, v in (line.split('=', maxsplit=1) for line in lines)} name = info['NAME'] if name.startswith("Debian"): system = System.Debian version = info['VERSION'].split()[0] # 11 (rodete) -> 11 elif name.startswith("CentOS"): system = System.CentOS version = info['VERSION_ID'] # 8 elif name.startswith("Rocky"): system = System.Rocky version = info['VERSION_ID'] # 8.4 elif name.startswith("Ubuntu"): system = System.Ubuntu version = info['VERSION_ID'] # 20.04 elif name.startswith("SLES"): system = System.SUSE version = info['VERSION_ID'] # 15.3 elif name.startswith("Red Hat"): system = System.RHEL version = info['VERSION_ID'] # 8.4 elif name.startswith("Fedora"): system = System.Fedora version = info['VERSION_ID'] # 34 else: raise RuntimeError("Unrecognized operating system.") return system, version def check_linux_distro(system: System, version: str) -> bool: """ Checks if given system version is supported by this script. Returns False if not, and prints information about the incompatibility. """ if '.' in version: version = version.split('.')[0] if len(SUPPORTED_SYSTEMS[system]) == 0: print_out(f"The {system} distribution is not supported by this script.") print_out("You may check https://docs.nvidia.com/cuda/cuda-installation-guide-linux/index.html") print_out("to try installing the drivers manually.") return False elif version not in SUPPORTED_SYSTEMS[system]: print_out(f"The version {version} of {system} is not supported by this script.") print_out(f"Supported versions: {SUPPORTED_SYSTEMS[system]}") print_out("You may try installing the driver manually following instructions from: ") print_out("https://docs.nvidia.com/cuda/cuda-installation-guide-linux/index.html") return False return True def check_driver_installed() -> bool: """ Checks if the driver is already installed by calling the `nvidia-smi` binary. If it's available, that means the driver is already installed. """ process = run("which nvidia-smi", check=False) if process.returncode != 0: return False process2 = run("nvidia-smi", check=False) return process2.returncode == 0 def install_dependencies_centos_rhel_rocky(system: System, version: str) -> bool: """ Installs required kernel-related packages and pciutils for CentOS and RHEL. """ if version.startswith("8"): binary = "dnf" else: binary = "yum" run(f"{binary} clean all") general_update = run(f"{binary} update -y --skip-broken") if "kernel" in general_update.stdout.decode(): return True # Kernel update requires a reboot before continuing kernel_install = run(f"{binary} install -y kernel") kernel_version = run("uname -r").stdout.decode().strip() if "already installed" not in kernel_install.stdout.decode(): return True # Kernel update requires a reboot if system == System.Rocky: if version.startswith("8"): run("dnf config-manager --set-enabled powertools") run(f"dnf config-manager --add-repo https://developer.download.nvidia.com/compute/cuda/repos/rhel{version[0]}/x86_64/cuda-rhel{version[0]}.repo") run("dnf update -y --skip-broken") run("dnf install -y epel-release") run(f"dnf install -y kernel-devel-{kernel_version} kernel-headers-{kernel_version}") elif system == System.CentOS and version.startswith("8"): run("dnf config-manager --set-enabled powertools") run("dnf install -y epel-release epel-next-release") elif system == System.CentOS and version.startswith("9"): run("dnf install -y https://dl.fedoraproject.org/pub/epel/next/9/Everything/x86_64/Packages/e/epel-next-release-9-1.el9.next.noarch.rpm") run("dnf install -y https://dl.fedoraproject.org/pub/epel/next/9/Everything/x86_64/Packages/e/epel-release-9-1.el9.next.noarch.rpm") elif system == System.RHEL and version[0] in ("8", "9"): run(f"dnf config-manager --add-repo https://developer.download.nvidia.com/compute/cuda/repos/rhel{version[0]}/x86_64/cuda-rhel{version[0]}.repo") run("dnf update -y --skip-broken") run(f"dnf install -y kernel-devel-{kernel_version} kernel-headers-{kernel_version}") run(f"dnf install -y https://dl.fedoraproject.org/pub/epel/epel-release-latest-{version[0]}.noarch.rpm") run(f"{binary} install -y kernel-devel epel-release " f"kernel-headers pciutils gcc make dkms acpid " f"libglvnd-glx libglvnd-opengl libglvnd-devel pkgconfig") return False def install_dependencies_sles(system: System, version: str) -> bool: # zypper install gcc make kernel-devel kernel-source # zypper install -t pattern devel_C_C++ devel_kernel # zypper install dkms return False # For now, there is not SLES script working. def install_dependencies_debian_ubuntu(system: System, version: str) -> bool: """ Prepares Debian and Ubuntu systems for driver installation. Makes sure the kernel is up to date, all required packages are installed and correct version of g++ will be used. """ # To make sure we don't get stuck waiting for user input. os.environ['DEBIAN_FRONTEND'] = 'noninteractive' kernel_version = run("uname -r").stdout.decode().strip() run("apt update") upgrade = run("apt upgrade -y").stdout.decode() if "Generating grub configuration file" in upgrade: # There was a kernel update, we need to reboot to work with proper kernel version return True run(f"apt install -y linux-headers-{kernel_version} " "software-properties-common pciutils gcc make dkms") if system == System.Ubuntu and version.startswith("22"): run("update-alternatives --install /usr/bin/gcc gcc /usr/bin/gcc-12 100 " "--slave /usr/bin/g++ g++ /usr/bin/g++-12 --slave /usr/bin/gcov gcov /usr/bin/gcov-12") return False def reboot(): """ Reboots the system. """ print_out("The system needs to be rebooted to complete the installation process. " "The process will be continued after the reboot.") print_out("Rebooting now.") run("reboot") sys.exit(0) def install_dependencies(system: System, version: str): """ Installs the driver dependencies to the system. This function may restart the system after installing some of the packages, in such situations the script should just be started again. """ if DEPENDENCIES_INSTALLED_FLAG.is_file(): return reboot_flag = False if system in (System.CentOS, System.RHEL, System.Rocky): reboot_flag = install_dependencies_centos_rhel_rocky(system, version) elif system in (System.Debian, System.Ubuntu): reboot_flag = install_dependencies_debian_ubuntu(system, version) elif system == System.SUSE: reboot_flag = install_dependencies_sles(system, version) else: raise RuntimeError("Unsupported operating system!") if reboot_flag: reboot() else: with DEPENDENCIES_INSTALLED_FLAG.open(mode='w') as flag: flag.write('1') if system == System.CentOS: # Both supported CentOS versions require reboot after this step reboot() def install_driver_runfile(system: System, version: str): dkms = "--dkms" if system in (System.RHEL, System.Rocky) and version.startswith("8"): # There is a problem with DKMS installation on Rocky and RHEL 8 # DKMS is disabled for those systems, you'll need to reinstall the drivers # with every kernel update. dkms = "" if detect_gpu_device() == TESLA_K80_DEVICE_CODE: run(f"curl -fSsl -O {K80_DRIVER_URL}") binary = f"NVIDIA-Linux-x86_64-{K80_DRIVER_VERSION}.run" else: run(f"curl -fSsl -O {DRIVER_URL}") binary = f"NVIDIA-Linux-x86_64-{DRIVER_VERSION}.run" attempt = 0 no_drm = "" while attempt < 3: install_run = run("sh {} -s {} {} --no-cc-version-check".format(binary, dkms, no_drm), check=False) if install_run.returncode == 0: return if "Failed to install the kernel module through DKMS" in install_run.stderr.decode(): dkms = "" if "--no-drm" in install_run.stderr.decode(): # Installer failed to install DRM KMS, so we try again with DRM disabled. no_drm = "--no-drm" attempt += 1 def post_install_steps(): """ Write the success message to log. """ with open(INSTALLER_DIR / 'success', mode='w') as success_file: success_file.write("Installation was completed on {}".format(datetime.now())) def parse_args(): parser = argparse.ArgumentParser(description='Install or verify GPU drivers.') parser.add_argument("action", choices=['install', 'verify'], default="install", nargs="?") parser.add_argument("--force", help="Forces driver installation, even if there is no GPU detected.", action='store_true') args = parser.parse_args() return args def install(args: argparse.Namespace): # Prerequisites check_python_version() if os.geteuid() != 0: print("This script needs to be run with root privileges!") sys.exit(1) # Set up the log directory. Logger.setup_log_dir() if check_driver_installed() and not args.force: print('Already installed.') sys.exit(0) # Check what system we're running system, version = detect_linux_distro() # Install the drivers and CUDA Toolkit install_dependencies(system, version) if not detect_gpu_device() and not args.force: print("There doesn't seem to be a GPU unit connected to your system. " "Aborting drivers installation.") sys.exit(0) install_driver_runfile(system, version) post_install_steps() def main(): """ Main function of the installation script. """ args = parse_args() if args.action == 'verify': if not check_driver_installed(): print("The driver is not installed.") else: print("The driver seems to be installed. Run `nvidia-smi` to check details.") return elif args.action == 'install': install(args) if __name__ == '__main__': try: main() except Exception as err: print_err("Failed with exception: " + str(err)) raise err