install.py

#!/usr/bin/env python

# Copyright 2021-2024 NVIDIA Corporation
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#

import argparse
import multiprocessing
import os
import platform
import shutil
import subprocess
import sys

# Flush output on newlines
sys.stdout.reconfigure(line_buffering=True)

os_name = platform.system()

if os_name == "Linux":
    pass
elif os_name == "Darwin":
    pass
else:
    raise Exception("install.py script does not work on %s" % os_name)


class BooleanFlag(argparse.Action):
    def __init__(
        self,
        option_strings,
        dest,
        default,
        required=False,
        help="",
        metavar=None,
    ):
        assert all(not opt.startswith("--no") for opt in option_strings)

        def flatten(list):
            return [item for sublist in list for item in sublist]

        option_strings = flatten(
            [
                (
                    [opt, "--no-" + opt[2:], "--no" + opt[2:]]
                    if opt.startswith("--")
                    else [opt]
                )
                for opt in option_strings
            ]
        )
        super().__init__(
            option_strings,
            dest,
            nargs=0,
            const=None,
            default=default,
            type=bool,
            choices=None,
            required=required,
            help=help,
            metavar=metavar,
        )

    def __call__(self, parser, namespace, values, option_string):
        setattr(namespace, self.dest, not option_string.startswith("--no"))


def execute_command(args, verbose, **kwargs):
    if verbose:
        print('Executing: "', " ".join(args), '" with ', kwargs)
    subprocess.check_call(args, **kwargs)


def scikit_build_cmake_build_dir(skbuild_dir):
    if os.path.exists(skbuild_dir):
        for f in os.listdir(skbuild_dir):
            if os.path.exists(
                cmake_build := os.path.join(skbuild_dir, f, "cmake-build")
            ):
                return cmake_build
    return None


def find_cmake_val(pattern, filepath):
    return (
        subprocess.check_output(["grep", "--color=never", pattern, filepath])
        .decode("UTF-8")
        .strip()
    )


def was_previously_built_with_different_build_isolation(
    isolated, legate_sparse_build_dir
):
    if (
        legate_sparse_build_dir is not None
        and os.path.exists(legate_sparse_build_dir)
        and os.path.exists(
            cmake_cache := os.path.join(legate_sparse_build_dir, "CMakeCache.txt")
        )
    ):
        try:
            if isolated:
                return True
            if find_cmake_val("pip-build-env", cmake_cache):
                return True
        except Exception:
            pass
    return False


def install_legate_sparse(
    arch,
    build_isolation,
    check_bounds,
    clean_first,
    cmake_exe,
    cmake_generator,
    conduit,
    cuda_dir,
    cuda,
    debug_release,
    debug,
    editable,
    extra_flags,
    gasnet_dir,
    networks,
    hdf,
    install_dir,
    legate_dir,
    llvm,
    march,
    maxdim,
    maxfields,
    nccl_dir,
    openmp,
    spy,
    thread_count,
    thrust_dir,
    unknown,
    verbose,
):
    if len(networks) > 1:
        print(
            "Warning: Building Realm with multiple networking backends is not "
            "fully supported currently."
        )

    if clean_first is None:
        clean_first = not editable

    print("Verbose build is ", "on" if verbose else "off")
    if verbose:
        print("Options are:")
        print("arch: ", arch)
        print("build_isolation: ", build_isolation)
        print("check_bounds: ", check_bounds)
        print("clean_first: ", clean_first)
        print("cmake_exe: ", cmake_exe)
        print("cmake_generator: ", cmake_generator)
        print("conduit: ", conduit)
        print("cuda_dir: ", cuda_dir)
        print("cuda: ", cuda)
        print("debug_release: ", debug_release)
        print("debug: ", debug)
        print("editable: ", editable)
        print("extra_flags: ", extra_flags)
        print("gasnet_dir: ", gasnet_dir)
        print("networks: ", networks)
        print("hdf: ", hdf)
        print("install_dir: ", install_dir)
        print("legate_dir: ", legate_dir)
        print("llvm: ", llvm)
        print("march: ", march)
        print("maxdim: ", maxdim)
        print("maxfields: ", maxfields)
        print("nccl_dir: ", nccl_dir)
        print("openmp: ", openmp)
        print("spy: ", spy)
        print("thread_count: ", thread_count)
        print("thrust_dir: ", thrust_dir)
        print("unknown: ", unknown)
        print("verbose: ", verbose)

    join = os.path.join
    exists = os.path.exists
    dirname = os.path.dirname
    realpath = os.path.realpath

    legate_sparse_dir = dirname(realpath(__file__))

    if thread_count is None:
        thread_count = multiprocessing.cpu_count()

    def validate_path(path):
        if path is not None and (path := str(path)) != "":
            if not os.path.isabs(path):
                path = join(legate_sparse_dir, path)
            if exists(path := realpath(path)):
                return path
        return None

    cuda_dir = validate_path(cuda_dir)
    nccl_dir = validate_path(nccl_dir)
    legate_dir = validate_path(legate_dir)
    thrust_dir = validate_path(thrust_dir)
    gasnet_dir = validate_path(gasnet_dir)

    if legate_dir is None:
        try:
            import legate.install_info as lg_install_info

            legate_dir = dirname(lg_install_info.libpath)
        except Exception:
            pass

    if verbose:
        print("cuda_dir: ", cuda_dir)
        print("nccl_dir: ", nccl_dir)
        print("legate_dir: ", legate_dir)
        print("thrust_dir: ", thrust_dir)
        print("gasnet_dir: ", gasnet_dir)

    skbuild_dir = join(legate_sparse_dir, "_skbuild")
    legate_sparse_build_dir = scikit_build_cmake_build_dir(skbuild_dir)

    if was_previously_built_with_different_build_isolation(
        build_isolation and not editable, legate_sparse_build_dir
    ):
        print("Performing a clean build to accommodate build isolation.")
        clean_first = True

    if clean_first:
        shutil.rmtree(skbuild_dir, ignore_errors=True)
        shutil.rmtree(join(legate_sparse_dir, "dist"), ignore_errors=True)
        shutil.rmtree(join(legate_sparse_dir, "build"), ignore_errors=True)
        shutil.rmtree(
            join(legate_sparse_dir, "legate-sparse.egg-info"),
            ignore_errors=True,
        )

    # Configure and build Legate Sparse via setup.py
    pip_install_cmd = [sys.executable, "-m", "pip", "install"]
    cmd_env = dict(os.environ.items())

    if unknown is not None:
        try:
            prefix_loc = unknown.index("--prefix")
            prefix_dir = validate_path(unknown[prefix_loc + 1])
            if prefix_dir is not None:
                install_dir = prefix_dir
                unknown = unknown[:prefix_loc] + unknown[prefix_loc + 2 :]
        except Exception:
            pass

    install_dir = validate_path(install_dir)

    if verbose:
        print("install_dir: ", install_dir)

    if install_dir is not None:
        pip_install_cmd += ["--root", "/", "--prefix", str(install_dir)]

    if editable:
        # editable implies build_isolation = False
        pip_install_cmd += ["--no-deps", "--no-build-isolation", "--editable"]
        cmd_env.update({"SETUPTOOLS_ENABLE_FEATURES": "legacy-editable"})
    else:
        if not build_isolation:
            pip_install_cmd += ["--no-deps", "--no-build-isolation"]
        pip_install_cmd += ["--upgrade"]

    if unknown is not None:
        pip_install_cmd += unknown

    pip_install_cmd += ["."]
    if verbose:
        pip_install_cmd += ["-vv"]

    # Also use preexisting CMAKE_ARGS from conda if set
    cmake_flags = cmd_env.get("CMAKE_ARGS", "").split(" ")

    if debug or verbose:
        cmake_flags += ["--log-level=%s" % ("DEBUG" if debug else "VERBOSE")]

    cmake_flags += f"""\
-DCMAKE_BUILD_TYPE={(
    "Debug" if debug else "RelWithDebInfo" if debug_release else "Release"
)}
-DBUILD_SHARED_LIBS=ON
-DBUILD_MARCH={str(march)}
-DCMAKE_CUDA_ARCHITECTURES={str(arch)}
-DLegion_MAX_DIM={str(maxdim)}
-DLegion_MAX_FIELDS={str(maxfields)}
-DLegion_SPY={("ON" if spy else "OFF")}
-DLegion_BOUNDS_CHECKS={("ON" if check_bounds else "OFF")}
-DLegion_USE_CUDA={("ON" if cuda else "OFF")}
-DLegion_USE_OpenMP={("ON" if openmp else "OFF")}
-DLegion_USE_LLVM={("ON" if llvm else "OFF")}
-DLegion_NETWORKS={";".join(networks)}
-DLegion_USE_HDF5={("ON" if hdf else "OFF")}
""".splitlines()

    if cuda_dir:
        cmake_flags += ["-DCUDAToolkit_ROOT=%s" % cuda_dir]
    if nccl_dir:
        cmake_flags += ["-DNCCL_DIR=%s" % nccl_dir]
    if gasnet_dir:
        cmake_flags += ["-DGASNet_ROOT_DIR=%s" % gasnet_dir]
    if conduit:
        cmake_flags += ["-DGASNet_CONDUIT=%s" % conduit]
    if thrust_dir:
        cmake_flags += ["-DThrust_ROOT=%s" % thrust_dir]
    if legate_dir:
        cmake_flags += ["-Dlegate_ROOT=%s" % legate_dir]

    cmake_flags += extra_flags
    build_flags = [f"-j{str(thread_count)}"]
    if verbose:
        if cmake_generator == "Unix Makefiles":
            build_flags += ["VERBOSE=1"]
        else:
            build_flags += ["--verbose"]

    cmd_env.update(
        {
            "CMAKE_ARGS": " ".join(cmake_flags),
            "CMAKE_GENERATOR": cmake_generator,
            "SKBUILD_BUILD_OPTIONS": " ".join(build_flags),
        }
    )

    execute_command(pip_install_cmd, verbose, cwd=legate_sparse_dir, env=cmd_env)


def driver():
    parser = argparse.ArgumentParser(description="Install Legate Sparse.")
    parser.add_argument(
        "--install-dir",
        dest="install_dir",
        metavar="DIR",
        required=False,
        default=None,
        help="Path to Legate Sparse source directory",
    )
    parser.add_argument(
        "--debug",
        dest="debug",
        action="store_true",
        required=False,
        default=os.environ.get("DEBUG", "0") == "1",
        help="Generates debug build",
    )
    parser.add_argument(
        "--debug-release",
        dest="debug_release",
        action="store_true",
        required=False,
        default=os.environ.get("DEBUG_RELEASE", "0") == "1",
        help="Generates release build with debugging symbols.",
    )
    parser.add_argument(
        "--check-bounds",
        dest="check_bounds",
        action="store_true",
        required=False,
        default=False,
        help="Build Legate Sparse with bounds checks.",
    )
    parser.add_argument(
        "--max-dim",
        dest="maxdim",
        type=int,
        default=int(os.environ.get("LEGION_MAX_DIM", 4)),
        help="Maximum number of dimensions supported",
    )
    parser.add_argument(
        "--max-fields",
        dest="maxfields",
        type=int,
        default=int(os.environ.get("LEGION_MAX_FIELDS", 256)),
        help="Maximum number of fields supported",
    )
    parser.add_argument(
        "--network",
        dest="networks",
        action="append",
        required=False,
        choices=["gasnet1", "gasnetex", "mpi"],
        default=[],
        help="Realm networking backend to use for multi-node execution.",
    )
    parser.add_argument(
        "--with-gasnet",
        dest="gasnet_dir",
        metavar="DIR",
        required=False,
        default=os.environ.get("GASNET"),
        help="Path to GASNet installation directory.",
    )
    parser.add_argument(
        "--with-legate",
        dest="legate_dir",
        metavar="DIR",
        required=False,
        default=os.environ.get("LEGATE_DIR"),
        help="Path to Legate installation directory.",
    )
    parser.add_argument(
        "--with-thrust",
        dest="thrust_dir",
        metavar="DIR",
        required=False,
        default=os.environ.get("THRUST_PATH"),
        help="Path to Thrust installation directory.",
    )
    parser.add_argument(
        "--with-nccl",
        dest="nccl_dir",
        metavar="DIR",
        required=False,
        default=os.environ.get("NCCL_PATH"),
        help="Path to NCCL installation directory.",
    )
    parser.add_argument(
        "--with-cmake",
        dest="cmake_exe",
        metavar="EXE",
        required=False,
        default="cmake",
        help="Path to CMake executable (if not on PATH).",
    )
    parser.add_argument(
        "--cmake-generator",
        dest="cmake_generator",
        required=False,
        default=os.environ.get(
            "CMAKE_GENERATOR",
            "Unix Makefiles" if shutil.which("ninja") is None else "Ninja",
        ),
        choices=["Ninja", "Unix Makefiles", None],
        help="The CMake makefiles generator",
    )
    parser.add_argument(
        "--cuda",
        action=BooleanFlag,
        default=os.environ.get("USE_CUDA", "0") == "1",
        help="Build with CUDA support to enable running on GPUs",
    )
    parser.add_argument(
        "--with-cuda",
        dest="cuda_dir",
        metavar="DIR",
        required=False,
        default=os.environ.get("CUDA"),
        help="Path to CUDA installation directory.",
    )
    parser.add_argument(
        "--arch",
        dest="arch",
        action="store",
        required=False,
        default="all-major",
        help="Specify the target GPU architecture.",
    )
    parser.add_argument(
        "--openmp",
        action=BooleanFlag,
        default=os.environ.get("USE_OPENMP", "0") == "1",
        help="Build with OpenMP support to enable using multiple CPU threads",
    )
    parser.add_argument(
        "--march",
        dest="march",
        required=False,
        default=("haswell" if platform.machine() == "x86_64" else "native"),
        help="Specify the target CPU architecture.",
    )
    parser.add_argument(
        "--llvm",
        dest="llvm",
        action="store_true",
        required=False,
        default=os.environ.get("USE_LLVM", "0") == "1",
        help="Build with LLVM support.",
    )
    parser.add_argument(
        "--hdf5",
        "--hdf",
        dest="hdf",
        action="store_true",
        required=False,
        default=os.environ.get("USE_HDF", "0") == "1",
        help="Build with HDF support.",
    )
    parser.add_argument(
        "--spy",
        dest="spy",
        action="store_true",
        required=False,
        default=os.environ.get("USE_SPY", "0") == "1",
        help="Build with detailed Legion Spy enabled.",
    )
    parser.add_argument(
        "--conduit",
        dest="conduit",
        action="store",
        required=False,
        # TODO: To support UDP conduit, we would need to add a special case on
        # the legate launcher.
        # See https://github.com/nv-legate/legate.core/issues/294.
        choices=["ibv", "ucx", "aries", "mpi"],
        default=os.environ.get("CONDUIT"),
        help="Build with specified GASNet conduit.",
    )
    parser.add_argument(
        "--clean",
        dest="clean_first",
        action=BooleanFlag,
        default=None,
        help="Clean before build.",
    )
    parser.add_argument(
        "--extra",
        dest="extra_flags",
        action="append",
        required=False,
        default=[],
        help="Extra CMake flags.",
    )
    parser.add_argument(
        "-j",
        dest="thread_count",
        nargs="?",
        type=int,
        required=False,
        default=os.environ.get("CPU_COUNT"),
        help="Number of threads used to compile.",
    )
    parser.add_argument(
        "--editable",
        dest="editable",
        action="store_true",
        required=False,
        default=False,
        help=(
            "Perform an editable install. Disables --build-isolation if set "
            "(passing --no-deps --no-build-isolation to pip)."
        ),
    )
    parser.add_argument(
        "--build-isolation",
        dest="build_isolation",
        action=BooleanFlag,
        required=False,
        default=True,
        help=(
            "Enable isolation when building a modern source distribution. "
            "Build dependencies specified by PEP 518 must be already "
            "installed if this option is used."
        ),
    )
    parser.add_argument(
        "-v",
        "--verbose",
        dest="verbose",
        action="store_true",
        required=False,
        default=False,
        help="Enable verbose build output.",
    )
    args, unknown = parser.parse_known_args()

    install_legate_sparse(unknown=unknown, **vars(args))


if __name__ == "__main__":
    driver()