# 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", "")