3131from hpccm .building_blocks .base import bb_base
3232from hpccm .building_blocks .generic_build import generic_build
3333from hpccm .building_blocks .packages import packages
34- from hpccm .common import linux_distro
35- from hpccm .config import get_cpu_architecture
34+ from hpccm .common import cpu_arch , linux_distro
3635from hpccm .primitives .comment import comment
3736from hpccm .primitives .copy import copy
3837from hpccm .primitives .environment import environment
@@ -56,7 +55,7 @@ class nccl(bb_base, hpccm.templates.downloader, hpccm.templates.envvars,
5655 use the latest commit on the default branch for the repository.
5756
5857 cuda: Flag to specify the CUDA version of the package to download.
59- The default is `11.6 `. This option is ignored if build is True.
58+ The default is `13.2 `. This option is ignored if build is True.
6059
6160 environment: Boolean flag to specify whether the environment
6261 (`CPATH`, `LD_LIBRARY_PATH`, `LIBRARY_PATH`, and `PATH`) should be
@@ -78,7 +77,7 @@ class nccl(bb_base, hpccm.templates.downloader, hpccm.templates.envvars,
7877 specified by `version`.
7978
8079 version: The version of NCCL to install. The default value is
81- `2.12.10 -1`.
80+ `2.29.7 -1`.
8281
8382 # Examples
8483
@@ -97,39 +96,47 @@ def __init__(self, **kwargs):
9796
9897 super (nccl , self ).__init__ (** kwargs )
9998
99+ self .__arch_label = '' # Filled in by __cpu_arch
100100 self .__baseurl = kwargs .pop ('baseurl' , 'https://github.com/NVIDIA/nccl/archive' )
101101 self .__build = kwargs .pop ('build' , False )
102102 self .__build_environment = '' # Filled in by __configure
103103 self .__default_repository = 'https://github.com/NVIDIA/nccl.git'
104104 self .__distro_label = '' # Filled in by __distro
105- self .__cuda = kwargs .pop ('cuda' , '11.6 ' )
105+ self .__cuda = kwargs .pop ('cuda' , '13.2 ' )
106106 self .__make_variables = kwargs .pop ('make_variables' , {})
107107 self .__ospackages = kwargs .pop ('ospackages' , [])
108108 self .__prefix = kwargs .pop ('prefix' , '/usr/local/nccl' )
109+ self .__repo_key = '' # Filled in by __repo_key
109110 self .__src_directory = kwargs .pop ('src_directory' , None )
110- self .__version = kwargs .pop ('version' , '2.12.10 -1' )
111+ self .__version = kwargs .pop ('version' , '2.29.7 -1' )
111112 self .__wd = kwargs .get ('wd' , hpccm .config .g_wd ) # working directory
112113
113114 if not self .__build :
114115 # Install prebuild package
115116
117+ # Set the CPU architecture specific parameters
118+ self .__cpu_arch ()
119+
116120 # Set the Linux distribution specific parameters
117121 self .__distro ()
118122
123+ # Set the repo key
124+ self .__repo ()
125+
119126 self += comment ('NCCL {}' .format (self .__version ))
120127 self += packages (ospackages = self .__ospackages )
121128 self += packages (
122129 apt = ['libnccl2={0}+cuda{1}' .format (self .__version ,
123130 self .__cuda ),
124131 'libnccl-dev={0}+cuda{1}' .format (self .__version ,
125132 self .__cuda )],
126- apt_keys = ['https://developer.download.nvidia.com/compute/cuda/repos/{0}/{1}/3bf863cc.pub ' .format (self .__distro_label , get_cpu_architecture () )],
127- apt_repositories = ['deb [signed-by=/usr/share/keyrings/3bf863cc.gpg ] https://developer.download.nvidia.com/compute/cuda/repos/{0}/{1} /' .format (self .__distro_label , get_cpu_architecture ( ))],
133+ apt_keys = ['https://developer.download.nvidia.com/compute/cuda/repos/{0}/{1}/{2} ' .format (self .__distro_label , self . __arch_label , self . __repo_key )],
134+ apt_repositories = ['deb [signed-by=/usr/share/keyrings/{2} ] https://developer.download.nvidia.com/compute/cuda/repos/{0}/{1} /' .format (self .__distro_label , self . __arch_label , self . __repo_key . replace ( '.pub' , '.gpg' ))],
128135 yum = ['libnccl-{0}+cuda{1}' .format (self .__version , self .__cuda ),
129136 'libnccl-devel-{0}+cuda{1}' .format (self .__version ,
130137 self .__cuda )],
131- yum_keys = ['https://developer.download.nvidia.com/compute/cuda/repos/{0}/{1}/3bf863cc.pub ' .format (self .__distro_label , get_cpu_architecture () )],
132- yum_repositories = ['https://developer.download.nvidia.com/compute/cuda/repos/{0}/{1}' .format (self .__distro_label , get_cpu_architecture () )])
138+ yum_keys = ['https://developer.download.nvidia.com/compute/cuda/repos/{0}/{1}/{2} ' .format (self .__distro_label , self . __arch_label , self . __repo_key )],
139+ yum_repositories = ['https://developer.download.nvidia.com/compute/cuda/repos/{0}/{1}' .format (self .__distro_label , self . __arch_label )])
133140
134141 else :
135142 # Build from source
@@ -166,6 +173,17 @@ def __init__(self, **kwargs):
166173 self += packages (ospackages = self .__ospackages )
167174 self += self .__bb
168175
176+ def __cpu_arch (self ):
177+ """Based on the CPU architecture, set values accordingly. A user
178+ specified value overrides any defaults."""
179+
180+ if hpccm .config .g_cpu_arch == cpu_arch .AARCH64 :
181+ self .__arch_label = 'sbsa'
182+ elif hpccm .config .g_cpu_arch == cpu_arch .X86_64 :
183+ self .__arch_label = 'x86_64'
184+ else : # pragma: no cover
185+ raise RuntimeError ('Unknown CPU architecture' )
186+
169187 def __configure (self ):
170188 """Setup build options based on user parameters"""
171189
@@ -192,16 +210,20 @@ def __distro(self):
192210 self .__ospackages = ['apt-transport-https' , 'ca-certificates' ,
193211 'gnupg' , 'wget' ]
194212
195- if hpccm .config .g_linux_version >= Version ('18.0' ):
196- self .__distro_label = 'ubuntu1804'
213+ if hpccm .config .g_linux_version >= Version ('24.0' ):
214+ self .__distro_label = 'ubuntu2404'
215+ elif hpccm .config .g_linux_version >= Version ('22.0' ):
216+ self .__distro_label = 'ubuntu2204'
197217 else :
198- self .__distro_label = 'ubuntu1604 '
218+ self .__distro_label = 'ubuntu2004 '
199219
200220 elif hpccm .config .g_linux_distro == linux_distro .CENTOS :
201- if hpccm .config .g_linux_version >= Version ('8.0' ):
202- self .__distro_label = 'rhel8'
221+ if hpccm .config .g_linux_version >= Version ('10.0' ):
222+ self .__distro_label = 'rhel10'
223+ elif hpccm .config .g_linux_version >= Version ('9.0' ):
224+ self .__distro_label = 'rhel9'
203225 else :
204- self .__distro_label = 'rhel7 '
226+ self .__distro_label = 'rhel8 '
205227
206228 else : # pragma: no cover
207229 raise RuntimeError ('Unknown Linux distribution' )
@@ -225,6 +247,21 @@ def __download(self):
225247 if not self .repository and not self .url :
226248 self .url = '{0}/v{1}.tar.gz' .format (self .__baseurl , self .__version )
227249
250+ def __repo (self ):
251+ """Based on the Linux distribution and CPU architecture, set values
252+ accordingly. A user specified value overrides any defaults.
253+ """
254+
255+ if self .__distro_label .startswith ('ubuntu' ):
256+ self .__repo_key = '3bf863cc.pub'
257+ elif self .__distro_label .startswith ('rhel' ):
258+ if hpccm .config .g_linux_version >= Version ('10.0' ):
259+ self .__repo_key = 'CDF6BA43.pub'
260+ else :
261+ self .__repo_key = 'D42D0685.pub'
262+ else : # pragma: no cover
263+ raise RuntimeError ('Unknown repository' )
264+
228265 def runtime (self , _from = '0' ):
229266 """Generate the set of instructions to install the runtime specific
230267 components from a build in a previous stage.
@@ -246,10 +283,10 @@ def runtime(self, _from='0'):
246283 self .rt += packages (
247284 apt = ['libnccl2={0}+cuda{1}' .format (self .__version ,
248285 self .__cuda )],
249- apt_keys = ['https://developer.download.nvidia.com/compute/cuda/repos/{0}/{1}/3bf863cc.pub ' .format (self .__distro_label , get_cpu_architecture () )],
250- apt_repositories = ['deb [signed-by=/usr/share/keyrings/3bf863cc.gpg ] https://developer.download.nvidia.com/compute/cuda/repos/{0}/{1} /' .format (self .__distro_label , get_cpu_architecture ( ))],
286+ apt_keys = ['https://developer.download.nvidia.com/compute/cuda/repos/{0}/{1}/{2} ' .format (self .__distro_label , self . __arch_label , self . __repo_key )],
287+ apt_repositories = ['deb [signed-by=/usr/share/keyrings/{2} ] https://developer.download.nvidia.com/compute/cuda/repos/{0}/{1} /' .format (self .__distro_label , self . __arch_label , self . __repo_key . replace ( '.pub' , '.gpg' ))],
251288 yum = ['libnccl-{0}+cuda{1}' .format (self .__version , self .__cuda )],
252- yum_keys = ['https://developer.download.nvidia.com/compute/cuda/repos/{0}/{1}/3bf863cc.pub ' .format (self .__distro_label , get_cpu_architecture () )],
253- yum_repositories = ['https://developer.download.nvidia.com/compute/cuda/repos/{0}/{1}' .format (self .__distro_label , get_cpu_architecture () )])
289+ yum_keys = ['https://developer.download.nvidia.com/compute/cuda/repos/{0}/{1}/{2} ' .format (self .__distro_label , self . __arch_label , self . __repo_key )],
290+ yum_repositories = ['https://developer.download.nvidia.com/compute/cuda/repos/{0}/{1}' .format (self .__distro_label , self . __arch_label )])
254291
255292 return str (self .rt )
0 commit comments