setup.py

# Copyright (c) EfficientMoE.
# SPDX-License-Identifier: Apache-2.0

# EfficientMoE Team

import io
import os
import sys

from setuptools import find_packages, setup

torch_available = True
try:
    import torch  # noqa: F401
except ImportError:
    torch_available = False
    print(
        "[WARNING] Unable to import torch, pre-compiling ops will be disabled. "
        "Please visit https://pytorch.org/ to see how to properly install torch on your system."
    )

ROOT_DIR = os.path.dirname(__file__)

sys.path.insert(0, ROOT_DIR)
# sys.path.insert(0, os.path.join(ROOT_DIR, 'src'))

from torch.utils import cpp_extension

from op_builder.all_ops import ALL_OPS

RED_START = "\033[31m"
RED_END = "\033[0m"
ERROR = f"{RED_START} [ERROR] {RED_END}"


def fetch_requirements(path):
    with open(path, "r") as fd:
        return [r.strip() for r in fd.readlines()]


def get_path(*filepath) -> str:
    return os.path.join(ROOT_DIR, *filepath)


def abort(msg):
    print(f"{ERROR} {msg}")
    assert False, msg


def read_readme() -> str:
    """Read the README file if present."""
    p = get_path("README.md")
    if os.path.isfile(p):
        return io.open(get_path("README.md"), "r", encoding="utf-8").read()
    else:
        return ""


install_requires = fetch_requirements("requirements.txt")

ext_modules = []

BUILD_OP_DEFAULT = int(os.environ.get("BUILD_OPS", 0))

if BUILD_OP_DEFAULT:
    assert torch_available, "Unable to pre-compile ops without torch installed. Please install torch before attempting to pre-compile ops."
    compatible_ops = dict.fromkeys(ALL_OPS.keys(), False)
    install_ops = dict.fromkeys(ALL_OPS.keys(), False)
    for op_name, builder in ALL_OPS.items():
        if builder is not None:
            op_compatible = builder.is_compatible()
            compatible_ops[op_name] = op_compatible
            if not op_compatible:
                abort(f"Unable to pre-compile {op_name}")
            ext_modules.append(builder.builder())

cmdclass = {
    "build_ext": cpp_extension.BuildExtension.with_options(use_ninja=True)
}

print(f"find_packages: {find_packages()}")

# install all files in the package, rather than just the egg
setup(
    name="moe_infinity",
    version=os.getenv("MOEINF_VERSION", "0.0.1"),
    packages=find_packages(
        exclude=["op_builder", "op_builder.*", "moe_infinity.ops.core.*"]
    ),
    package_data={
        "moe_infinity.ops.prefetch": ["**/*.so"],
        "moe_infinity": ["ops/core/**"],
    },
    include_package_data=True,
    install_requires=install_requires,
    author="EfficientMoE Team",
    long_description=read_readme(),
    long_description_content_type="text/markdown",
    url="https://github.com/EfficientMoE/MoE-Infinity",
    project_urls={"Homepage": "https://github.com/EfficientMoE/MoE-Infinity"},
    classifiers=[
        "Programming Language :: Python :: 3.8",
        "Programming Language :: Python :: 3.9",
        "Programming Language :: Python :: 3.10",
        "Programming Language :: Python :: 3.11",
        "License :: OSI Approved :: Apache Software License",
        "Topic :: Scientific/Engineering :: Artificial Intelligence",
    ],
    license="Apache License 2.0",
    python_requires=">=3.8",
    ext_modules=ext_modules,
    cmdclass=cmdclass,
)