diff options
| -rw-r--r-- | var/spack/repos/builtin/packages/nccl-fastsocket/package.py | 65 |
1 files changed, 65 insertions, 0 deletions
diff --git a/var/spack/repos/builtin/packages/nccl-fastsocket/package.py b/var/spack/repos/builtin/packages/nccl-fastsocket/package.py new file mode 100644 index 0000000000..f25862b2eb --- /dev/null +++ b/var/spack/repos/builtin/packages/nccl-fastsocket/package.py @@ -0,0 +1,65 @@ +# Copyright 2013-2023 Lawrence Livermore National Security, LLC and other +# Spack Project Developers. See the top-level COPYRIGHT file for details. +# +# SPDX-License-Identifier: (Apache-2.0 OR MIT) + +import tempfile + +from spack.package import * + + +class NcclFastsocket(Package): + """NCCL Fast Socket GCP Net plugin for NCCL""" + + homepage = "https://github.com/google/nccl-fastsocket" + git = "https://github.com/google/nccl-fastsocket.git" + + version("master", preferred=True) + + depends_on("bazel", type="build") + depends_on("nccl", type=["build", "run"]) + + maintainers("danielahlin") + + def setup_build_environment(self, env): + spec = self.spec + tmp_path = tempfile.mkdtemp(prefix="spack") + env.set("TEST_TMPDIR", tmp_path) + env.set("NCCL_INSTALL_PATH", spec["nccl"].prefix) + env.set("NCCL_HDR_PATH", spec["nccl"].prefix.include) + + def install(self, spec, prefix): + tmp_path = env["TEST_TMPDIR"] + # Copied of py-tensorflow + args = [ + # Don't allow user or system .bazelrc to override build settings + "--nohome_rc", + "--nosystem_rc", + # Bazel does not work properly on NFS, switch to /tmp + "--output_user_root=" + tmp_path, + "build", + "libnccl-net.so", + # Spack logs don't handle colored output well + "--color=no", + "--jobs={0}".format(make_jobs), + # Enable verbose output for failures + "--verbose_failures", + # Show (formatted) subcommands being executed + "--subcommands=pretty_print", + # Ask bazel to explain what it's up to + # Needs a filename as argument + "--explain=explainlogfile.txt", + # Increase verbosity of explanation, + "--verbose_explanations", + ] + bazel(*args) + install_tree("bazel-bin", prefix.lib) + + def setup_run_environment(self, env): + # The current plugin pickup method of NCCL is to scan for libraries with certain + # names in the standard library search paths. Consequently, to make nccl-fastsocket + # discoverable to NCCL it is necessary to add it to the LD_LIBRARY_PATH. + env.prepend_path("LD_LIBRARY_PATH", self.prefix.lib) + # NCCL_NET_PLUGIN can be used to change part of the library-name NCCL is looking + # for. This is not necessary for this plugin so it is set to the empty string. + env.set("NCCL_NET_PLUGIN", "") |