# Copyright 2013-2024 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 os
from spack.package import *
class RedisAi(MakefilePackage):
"""A Redis module for serving tensors and executing deep learning graphs"""
homepage = "https://oss.redis.com/redisai/"
git = "https://github.com/RedisAI/RedisAI.git"
maintainers("MattToast")
license("Apache-2.0")
version(
"1.2.7", tag="v1.2.7", commit="1bf38d86233ba06e1350ca9de794df2b07cdb274", submodules=True
)
variant("torch", default=True, description="Build with the pytorch backend")
variant("cuda", default=False, description="Use CUDA")
variant("rocm", default=False, description="Use ROCm")
conflicts("+cuda+rocm")
# Required dependencies
depends_on("git", type=("build", "link"))
depends_on("git-lfs", type=("build", "link"))
depends_on("python@3:", type=("build", "link"))
depends_on("py-pip", type=("build", "link"))
depends_on("cmake@3.0:", type=("build", "link"))
depends_on("gmake", type=("build", "link"))
# GPU deps
depends_on("cuda@11.2:", type=("build", "link", "run"), when="+cuda")
depends_on("cudnn@8.1:", type=("build", "link", "run"), when="+cuda")
with when("+rocm"):
depends_on("hsa-rocr-dev")
depends_on("hip")
depends_on("rocprim")
depends_on("hipcub")
depends_on("rocthrust")
depends_on("roctracer-dev")
depends_on("rocrand")
depends_on("hipsparse")
depends_on("hipfft")
depends_on("rocfft")
depends_on("rocblas")
depends_on("miopen-hip")
depends_on("rocminfo")
# Optional Deps
with when("+torch"):
depends_on("py-torch@1.11.0:~cuda~rocm", type=("build", "link"), when="~cuda~rocm")
depends_on("py-torch@1.11.0:+cuda+cudnn~rocm", type=("build", "link"), when="+cuda")
depends_on("py-torch@1.11.0:~cuda+rocm", type=("build", "link"), when="+rocm")
build_directory = "opt"
parallel = False
@property
def use_gpu(self):
return self.spec.satisfies("+cuda") or self.spec.satisfies("+rocm")
@property
def with_torch(self):
return self.spec.satisfies("+torch")
@property
def torch_dir(self):
return (
join_path(self.spec["py-torch"].package.cmake_prefix_paths[0], "Torch")
if self.with_torch
else None
)
@property
def build_env(self):
build_env = {
"WITH_TF": "0",
"WITH_TFLITE": "0",
"WITH_PT": "0",
"WITH_ORT": "0",
"WITH_UNIT_TESTS": "0",
"GPU": "1" if self.use_gpu else "0",
}
if self.with_torch:
build_env["WITH_PT"] = "1"
build_env["Torch_DIR"] = self.torch_dir
return build_env
def edit(self, spec, prefix):
# resolve deps not provided through spack
Executable(join_path(".", "get_deps.sh"))(
extra_env={
"VERBOSE": "1",
# Need to grab the RAI specific version of dlpack
"WITH_DLPACK": "1",
# Do not get ml backends, they should be retrieved through spack
"WITH_TF": "0",
"WITH_TFLITE": "0",
"WITH_PT": "0",
"WITH_ORT": "0",
# Decide if we want GPU
"GPU": "1" if self.use_gpu else "0",
}
)
env.update(self.build_env)
def install(self, spec, prefix):
super(RedisAi, self).install(spec, prefix)
install_tree("install-*", prefix)
@run_after("install")
@on_package_attributes(with_torch=True)
def copy_libtorch(self):
torch_site_dir = os.path.dirname(os.path.dirname(os.path.dirname(self.torch_dir)))
torch_lib_dir = join_path(torch_site_dir, "lib")
install_tree(torch_lib_dir, self.prefix.backends.redisai_torch.lib)
def setup_run_environment(self, env):
env.set("REDIS_AI", self.prefix.join("redisai.so"))