cbensimon HF Staff commited on
Commit
d227777
·
1 Parent(s): 2e8adf4

CUDA toolkit install script

Browse files
Files changed (2) hide show
  1. utils/__init__.py +0 -0
  2. utils/cuda_toolkit.py +32 -0
utils/__init__.py ADDED
File without changes
utils/cuda_toolkit.py ADDED
@@ -0,0 +1,32 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """
2
+ """
3
+
4
+ import os
5
+ import subprocess
6
+ from pathlib import Path
7
+
8
+
9
+ CUDA_VERSION = '12.4.0'
10
+ DRIVERS_VERSION = '550.54.14'
11
+
12
+
13
+ def install_cuda_toolkit():
14
+ installer_name = f'cuda_{CUDA_VERSION}_{DRIVERS_VERSION}_linux.run'
15
+ installer_url = f'https://developer.download.nvidia.com/compute/cuda/{CUDA_VERSION}/local_installers/{installer_name}'
16
+ installer_path = f'/tmp/{installer_name}'
17
+ cuda_path = f'/usr/local/cuda-{CUDA_VERSION}'
18
+ bashrc_path = Path.home() / '.bashrc'
19
+
20
+ subprocess.run('wget -q --show-progress {installer_url} -O {installer_path}', check=True, shell=True)
21
+ subprocess.run('fakeroot sh {installer_path} --silent --toolkit --override', check=True, shell=True)
22
+
23
+ bashrc_lines = (
24
+ f'# CUDA toolkit',
25
+ f'export PATH={cuda_path}/bin:$PATH',
26
+ f'export LD_LIBRARY_PATH={cuda_path}/lib64:$LD_LIBRARY_PATH'
27
+ )
28
+ with open(bashrc_path, 'a') as f:
29
+ f.write('\n' + '\n'.join(bashrc_lines) + '\n')
30
+
31
+ os.environ['PATH'] = f"{cuda_path}/bin:{os.environ.get('PATH', '')}"
32
+ os.environ['LD_LIBRARY_PATH'] = f"{cuda_path}/lib64:{os.environ.get('LD_LIBRARY_PATH', '')}"