Updated NCCL package installer to take in CUDA related parameters and…

… to work with existing NCCL download ------------- Created by MOE: https://github.com/google/moe MOE_MIGRATED_REVID=284262515
antigenius0910 · Dec 11, 2019 · 124167c · 124167c
1 parent 58b89c5
commit 124167c
Show file tree

Hide file tree

Showing 2 changed files with 40 additions and 8 deletions.
diff --git a/perfkitbenchmarker/linux_packages/nccl.py b/perfkitbenchmarker/linux_packages/nccl.py
@@ -14,28 +14,44 @@
 
 
 """Module containing NCCL installation function."""
-
+import posixpath
 from perfkitbenchmarker import flags
+from perfkitbenchmarker import vm_util
 
 flags.DEFINE_string('nccl_version', '2.5.6-1',
                     'NCCL version to install')
 
 FLAGS = flags.FLAGS
 
+GIT_REPO = 'https://github.com/NVIDIA/nccl.git'
+
 
 def _Build(vm):
   """Installs the OpenMPI package on the VM."""
-  vm.RemoteCommand('git clone https://github.com/NVIDIA/nccl.git --branch '
-                   'v{}'.format(FLAGS.nccl_version))
-  vm.RemoteCommand('cd nccl && make -j src.build')
+  vm.RemoteCommand('[ -d "nccl" ] || git clone {git_repo} --branch v{version}'
+                   .format(git_repo=GIT_REPO, version=FLAGS.nccl_version))
+  cuda_home = '/usr/local/cuda'
+  vm.InstallPackages('build-essential devscripts debhelper fakeroot')
+
+  env_vars = {}
+  env_vars['PATH'] = (r'{cuda_bin_path}:$PATH'
+                      .format(cuda_bin_path=posixpath.join(cuda_home, 'bin')))
+  env_vars['CUDA_HOME'] = (r'{cuda_home}'.format(cuda_home=cuda_home))
+  env_vars['LD_LIBRARY_PATH'] = (r'{lib_path}:$LD_LIBRARY_PATH'
+                                 .format(lib_path=posixpath.join(
+                                     cuda_home, 'lib64')))
+
+  vm.RemoteCommand('cd nccl && {env} make -j 20 pkg.debian.build'
+                   .format(env=vm_util.DictonaryToEnvString(env_vars)))
 
 
 def AptInstall(vm):
   """Installs the NCCL package on the VM."""
   _Build(vm)
-  vm.InstallPackages('build-essential devscripts debhelper fakeroot')
-  vm.RemoteCommand('cd nccl && make pkg.debian.build')
+
   vm.InstallPackages('{build}libnccl2_{nccl}+cuda{cuda}_amd64.deb '
-                     '{build}libnccl-dev_{nccl}+cuda{cuda}_amd64.deb'.format(
-                         build='./nccl/build/pkg/deb/', nccl=FLAGS.nccl_version,
+                     '{build}libnccl-dev_{nccl}+cuda{cuda}_amd64.deb'
+                     .format(
+                         build='./nccl/build/pkg/deb/',
+                         nccl=FLAGS.nccl_version,
                          cuda=FLAGS.cuda_toolkit_version))
diff --git a/perfkitbenchmarker/vm_util.py b/perfkitbenchmarker/vm_util.py
@@ -671,3 +671,19 @@ def ReplaceText(vm, current_value, new_value, file_name, regex_char='/'):
                        current_value=current_value,
                        new_value=new_value,
                        file=file_name))
+
+
+def DictonaryToEnvString(dictionary):
+  """Convert a dictionary to a space sperated 'key=value' string.
+
+  Args:
+    dictionary: the key-value dictionary to be convert
+
+  Returns:
+    a string representing the dictionary
+  """
+  dict_str = ''
+  for key, value in sorted(dictionary.items()):
+    dict_str += ' {key}={value}'.format(key=key, value=value)
+  return dict_str
+