diff options
Diffstat (limited to 'var/spack/repos/builtin/packages/redis-ai/package.py')
-rw-r--r-- | var/spack/repos/builtin/packages/redis-ai/package.py | 124 |
1 files changed, 124 insertions, 0 deletions
diff --git a/var/spack/repos/builtin/packages/redis-ai/package.py b/var/spack/repos/builtin/packages/redis-ai/package.py new file mode 100644 index 0000000000..e54b10fb5e --- /dev/null +++ b/var/spack/repos/builtin/packages/redis-ai/package.py @@ -0,0 +1,124 @@ +# 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 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") + + version("1.2.7", tag="v1.2.7", 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")) |