""" """ import os import subprocess from pathlib import Path CUDA_VERSION = '12.4.0' DRIVERS_VERSION = '550.54.14' def install_cuda_toolkit(): installer_name = f'cuda_{CUDA_VERSION}_{DRIVERS_VERSION}_linux.run' installer_url = f'https://developer.download.nvidia.com/compute/cuda/{CUDA_VERSION}/local_installers/{installer_name}' installer_path = f'/tmp/{installer_name}' cuda_path = f'/usr/local/cuda-{CUDA_VERSION}' bashrc_path = Path.home() / '.bashrc' subprocess.run(f'wget --progress=dot:giga {installer_url} -O {installer_path}', check=True, shell=True) subprocess.run(f'fakeroot sh {installer_path} --silent --toolkit --override', check=True, shell=True) bashrc_lines = ( f'# CUDA toolkit', f'export PATH={cuda_path}/bin:$PATH', f'export LD_LIBRARY_PATH={cuda_path}/lib64:$LD_LIBRARY_PATH' ) with open(bashrc_path, 'a') as f: f.write('\n' + '\n'.join(bashrc_lines) + '\n') os.environ['PATH'] = f"{cuda_path}/bin:{os.environ.get('PATH', '')}" os.environ['LD_LIBRARY_PATH'] = f"{cuda_path}/lib64:{os.environ.get('LD_LIBRARY_PATH', '')}"