setup.py

# Copyright © Advanced Micro Devices, Inc. All rights reserved.
#
# MIT License
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in all
# copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.
import os
import re
import subprocess
import sys
from pathlib import Path
import shutil

from setuptools import Extension, find_packages, setup
from setuptools.command.build import build as _build
from setuptools.command.build_ext import build_ext

try:
    from torch.utils.cpp_extension import CppExtension as _TorchCppExtension
except ImportError:  # torch is optional at build time
    _TorchCppExtension = None

try:
    from Cython.Build import cythonize as _cythonize

    _HAVE_CYTHON = True
except ImportError:
    _HAVE_CYTHON = False

_supported_arch_list = ["gfx942", "gfx950", "gfx1250"]

_REQUIRED_SYSTEM_DEPS: list = []

_MPI_SYSTEM_DEPS = [
    (
        "mpicc",
        ("libopenmpi-dev", "openmpi-devel"),
        "MPI compiler wrapper (needed by CMake)",
    ),
    ("mpirun", ("openmpi-bin", "openmpi"), "MPI runtime (needed at runtime)"),
]

_REQUIRED_HEADERS = [
    (
        ["/usr/include/infiniband/verbs.h"],
        ("libibverbs-dev", "rdma-core-devel"),
        "InfiniBand verbs headers (needed for RDMA transport)",
    ),
]


def _env_flag(name: str, default: str = "OFF") -> bool:
    return os.environ.get(name, default).strip().upper() in {
        "1",
        "ON",
        "TRUE",
        "YES",
    }


def _torch_symm_mode() -> str:
    """ON | OFF | AUTO for the torch SymmetricMemory backend.

    AUTO (the default) builds it when torch is importable and treats a failure as a
    skip, because the extension is optional and its ABI tracks a torch that mori does
    not depend on. ON turns those skips into build errors."""
    raw = os.environ.get("BUILD_TORCH_SYMM", "AUTO").strip().upper()
    if raw in {"1", "ON", "TRUE", "YES"}:
        return "ON"
    if raw in {"0", "OFF", "FALSE", "NO"}:
        return "OFF"
    return "AUTO"


def _detect_pkg_manager() -> str:
    """Detect the system package manager."""
    if shutil.which("apt-get"):
        return "apt"
    if shutil.which("dnf"):
        return "dnf"
    if shutil.which("yum"):
        return "yum"
    return "unknown"


def _check_system_deps() -> None:
    """Verify required system packages are installed; print install hints if not."""
    missing = []

    for binary, pkgs, desc in _REQUIRED_SYSTEM_DEPS:
        if not shutil.which(binary):
            missing.append((pkgs, desc))

    for paths, pkgs, desc in _REQUIRED_HEADERS:
        if not any(os.path.isfile(p) for p in paths):
            missing.append((pkgs, desc))

    if not missing:
        return

    pm = _detect_pkg_manager()
    pkg_idx = 0 if pm == "apt" else 1

    lines = ["", "=" * 70, "[mori] Missing system dependencies:"]
    for pkgs, desc in missing:
        pkg_name = pkgs[pkg_idx] if isinstance(pkgs, tuple) else pkgs
        lines.append(f"  - {pkg_name:24s}  {desc}")
    lines.append("")

    pkg_names = [(p[pkg_idx] if isinstance(p, tuple) else p) for p, _ in missing]
    if pm == "apt":
        lines.append("  Install (Ubuntu/Debian):")
        lines.append(
            f"    sudo apt-get update && sudo apt-get install -y {' '.join(pkg_names)}"
        )
    elif pm in ("dnf", "yum"):
        lines.append("  Install (RHEL/CentOS/Fedora):")
        lines.append(f"    sudo {pm} install -y {' '.join(pkg_names)}")
    else:
        lines.append("  Install the equivalent packages for your distribution:")
        lines.append(
            f"    Ubuntu/Debian: sudo apt-get install {' '.join(p[0] if isinstance(p, tuple) else p for p, _ in missing)}"
        )
        lines.append(
            f"    RHEL/Fedora:   sudo dnf install {' '.join(p[1] if isinstance(p, tuple) else p for p, _ in missing)}"
        )
    lines.append("=" * 70)
    print("\n".join(lines), file=sys.stderr)
    raise RuntimeError(
        f"Missing system packages: {', '.join(pkg_names)}. "
        "See messages above for install instructions."
    )


# System packages needed only by UMBP's distributed control plane (gRPC +
# Protobuf). These are apt/dnf packages, not pip-installable, so we detect them
# and gracefully disable UMBP when absent instead of failing the whole build.
_UMBP_APT_PACKAGES = [
    "libgrpc-dev",
    "libgrpc++-dev",
    "libprotobuf-dev",
    "protobuf-compiler-grpc",
]


def _detect_missing_umbp_deps() -> list:
    """Return the UMBP gRPC/Protobuf build deps that appear to be missing.

    Mirrors what src/umbp/CMakeLists.txt requires: the C++ gRPC/Protobuf headers
    plus the protoc compiler and the gRPC C++ plugin.
    """
    missing = []
    header_checks = [
        (
            [
                "/usr/include/grpcpp/grpcpp.h",
                "/usr/include/x86_64-linux-gnu/grpcpp/grpcpp.h",
            ],
            "libgrpc++-dev",
        ),
        (
            [
                "/usr/include/google/protobuf/message.h",
                "/usr/include/x86_64-linux-gnu/google/protobuf/message.h",
            ],
            "libprotobuf-dev",
        ),
    ]
    for paths, pkg in header_checks:
        if not any(os.path.isfile(p) for p in paths):
            missing.append(pkg)
    if not shutil.which("protoc"):
        missing.append("protobuf-compiler")
    if not shutil.which("grpc_cpp_plugin"):
        missing.append("protobuf-compiler-grpc")
    return missing


def _warn_umbp_disabled(missing: list, explicit: bool) -> None:
    """Warn that UMBP is being disabled because gRPC/Protobuf are unavailable."""
    pm = _detect_pkg_manager()
    lines = [
        "",
        "=" * 70,
        "[mori] UMBP needs gRPC + Protobuf, but these system packages appear "
        "to be missing:",
    ]
    for pkg in missing:
        lines.append(f"  - {pkg}")
    lines.append("")
    if explicit:
        lines.append(
            "  You requested UMBP (BUILD_UMBP / BUILD_UMBP_SPDK), but it will be "
            "DISABLED for this build."
        )
    else:
        lines.append("  UMBP will be DISABLED for this build.")
    lines.append("  All other mori modules will still be built and installed.")
    lines.append("")
    lines.append("  To build UMBP, install the dependencies and rebuild:")
    if pm in ("dnf", "yum"):
        lines.append(
            f"    sudo {pm} install -y grpc-devel grpc-plugins "
            "protobuf-devel protobuf-compiler"
        )
    else:
        lines.append(
            "    sudo apt-get update && sudo apt-get install -y "
            + " ".join(_UMBP_APT_PACKAGES)
        )
    lines.append("=" * 70)
    print("\n".join(lines), file=sys.stderr)


def _invalidate_cmake_cache_if_changed(cmake_cache: "Path", cmake_args: list) -> None:
    """Clear CMake cache if any -DKEY=VALUE arg differs from the cached value."""
    if not cmake_cache.is_file():
        return

    # Parse -DKEY=VALUE args (normalize booleans to uppercase)
    _BOOL_MAP = {
        "1": "ON",
        "TRUE": "ON",
        "YES": "ON",
        "0": "OFF",
        "FALSE": "OFF",
        "NO": "OFF",
    }

    def _normalize(v: str) -> str:
        return _BOOL_MAP.get(v.upper(), v)

    new_opts: dict[str, str] = {}
    for arg in cmake_args:
        if arg.startswith("-D") and "=" in arg:
            key, val = arg[2:].split("=", 1)
            new_opts[key] = _normalize(val)

    # Parse CMakeCache.txt: lines like KEY:TYPE=VALUE
    cached_opts: dict[str, str] = {}
    for line in cmake_cache.read_text().splitlines():
        line = line.strip()
        if line.startswith("#") or line.startswith("//") or "=" not in line:
            continue
        key_type, val = line.split("=", 1)
        key = key_type.split(":")[0]
        cached_opts[key] = _normalize(val)

    changed = [
        k for k, v in new_opts.items() if k in cached_opts and cached_opts[k] != v
    ]

    # Also check stale CMAKE_MAKE_PROGRAM path
    make_prog = cached_opts.get("CMAKE_MAKE_PROGRAM", "")
    if make_prog and not os.path.isfile(make_prog):
        changed.append("CMAKE_MAKE_PROGRAM (no longer exists)")

    if changed:
        print(f"[mori] CMake options changed ({', '.join(changed)}), clearing cache.")
        cmake_cache.unlink()
        cmake_files = cmake_cache.parent / "CMakeFiles"
        if cmake_files.is_dir():
            shutil.rmtree(cmake_files)


def _detect_local_gpu_arch() -> str | None:
    """Auto-detect the GPU architecture on the current machine."""
    rocm_path = os.environ.get("ROCM_PATH", "/opt/rocm")
    enumerator = os.path.join(rocm_path, "bin", "rocm_agent_enumerator")
    if os.path.isfile(enumerator):
        try:
            out = subprocess.check_output([enumerator], text=True)
            for line in out.strip().split("\n"):
                line = line.strip()
                if (
                    line.startswith("gfx")
                    and line != "gfx000"
                    and line in _supported_arch_list
                ):
                    return line
        except subprocess.CalledProcessError:
            pass
    return None


def _get_gpu_archs() -> str:
    """Determine GPU target architectures for compilation.

    Priority: MORI_GPU_ARCHS > local GPU > PYTORCH_ROCM_ARCH / GPU_ARCHS > fat binary default.
    """
    mori_gpu_archs = os.environ.get("MORI_GPU_ARCHS", None)
    if mori_gpu_archs:
        return mori_gpu_archs

    local_arch = _detect_local_gpu_arch()
    if local_arch:
        return local_arch

    archs = os.environ.get("PYTORCH_ROCM_ARCH", None)

    gpu_archs = os.environ.get("GPU_ARCHS", None)
    if gpu_archs:
        archs = gpu_archs

    if archs:
        arch_list = archs.replace(" ", ";").split(";")
        valid_arch_list = list(set(_supported_arch_list) & set(arch_list))
        if valid_arch_list:
            return ";".join(valid_arch_list)

    print(
        f"[mori] No GPU arch specified — building fat binary for {_supported_arch_list}"
    )
    return ";".join(_supported_arch_list)


def _copy_jit_sources(root_dir: Path) -> None:
    """Copy JIT-required source files into the package for wheel distribution.

    This creates python/mori/_jit-sources/ with the same directory structure
    as the repo root, so that get_mori_source_root() can use it as a drop-in
    replacement when the original source tree is not available.
    """
    jit_dir = root_dir / "python" / "mori" / "_jit-sources"
    if jit_dir.exists():
        shutil.rmtree(jit_dir)

    def _copytree(src, dst, **kw):
        shutil.copytree(src, dst, dirs_exist_ok=True, **kw)

    _copytree(root_dir / "include", jit_dir / "include")

    _copytree(root_dir / "src" / "ops" / "kernels", jit_dir / "src" / "ops" / "kernels")
    _copytree(
        root_dir / "src" / "ops" / "dispatch_combine",
        jit_dir / "src" / "ops" / "dispatch_combine",
    )
    # JIT v2 device bodies. The rendered TU includes one of these by path, so
    # without them the v2 hip backend does not compile at all -- and only off a
    # repo checkout would anyone notice.
    _copytree(
        root_dir / "src" / "ops" / "dispatch_combine_v2",
        jit_dir / "src" / "ops" / "dispatch_combine_v2",
    )

    io_kernels_src = root_dir / "src" / "io" / "kernels"
    if io_kernels_src.is_dir():
        _copytree(io_kernels_src, jit_dir / "src" / "io" / "kernels")

    ccl_kernels_src = root_dir / "src" / "collective" / "kernels"
    if ccl_kernels_src.is_dir():
        _copytree(ccl_kernels_src, jit_dir / "src" / "collective" / "kernels")

    shmem_dst = jit_dir / "src" / "shmem"
    shmem_dst.mkdir(parents=True, exist_ok=True)
    for name in ["shmem_device_api_wrapper.cpp"]:
        src_file = root_dir / "src" / "shmem" / name
        if src_file.is_file():
            shutil.copy2(src_file, shmem_dst / name)

    # torch SymmetricMemory backend — compiled on first import when the prebuilt
    # extension is missing or was built against a different torch. Headers come from
    # the include/ copy above.
    symm_src = root_dir / "src" / "allocator" / "symm_backend.cpp"
    if symm_src.is_file():
        symm_dst = jit_dir / "src" / "allocator"
        symm_dst.mkdir(parents=True, exist_ok=True)
        shutil.copy2(symm_src, symm_dst / "symm_backend.cpp")

    # cco device-API wrapper — JIT-compiled to libmori_cco_device.bc on first use
    # (mori.cco.device.bitcode). Headers come from the include/ copy above.
    cco_dev_src = root_dir / "src" / "cco" / "device" / "cco_device_wrapper.cpp"
    if cco_dev_src.is_file():
        cco_dst = jit_dir / "src" / "cco" / "device"
        cco_dst.mkdir(parents=True, exist_ok=True)
        shutil.copy2(cco_dev_src, cco_dst / "cco_device_wrapper.cpp")

    for subdir in ["spdlog/include", "msgpack-c/include"]:
        src = root_dir / "3rdparty" / subdir
        if src.is_dir():
            _copytree(src, jit_dir / "3rdparty" / subdir)

    profiler_tools_src = root_dir / "tools" / "profiler"
    if profiler_tools_src.is_dir():
        _copytree(profiler_tools_src, jit_dir / "tools" / "profiler")


def _write_cco_build_flags(root_dir: Path, *, sdma: bool) -> None:
    """Bake the SDMA build flag into the package so the runtime JIT matches it.

    Writes python/mori/cco/device/_build_flags.py; mori.cco.device.bitcode reads
    it so `BUILD_CCO_SDMA=ON pip install .` JITs a matching device wrapper by
    default (env still overrides).
    """
    flags_path = root_dir / "python" / "mori" / "cco" / "device" / "_build_flags.py"
    flags_path.parent.mkdir(parents=True, exist_ok=True)
    flags_path.write_text(
        "# Auto-generated by setup.py at build time — do not edit.\n"
        f"BUILD_CCO_SDMA = {bool(sdma)!r}\n"
    )


_ROCM_VERSION_RE = re.compile(r"\b(\d+)\.(\d+)(?:\.\d+)?\b")


def _detect_rocm_build_version() -> tuple[str | None, str | None]:
    """Return the build ROCm version and a major.minor compatibility pattern."""

    rocm_path = Path(
        os.environ.get("ROCM_PATH") or os.environ.get("HIP_PATH") or "/opt/rocm"
    ).resolve()
    candidates: list[str] = []

    for version_file in (
        rocm_path / ".info" / "version",
        rocm_path / ".info" / "version-dev",
    ):
        try:
            candidates.append(version_file.read_text(encoding="utf-8"))
        except OSError:
            pass

    hipconfig = rocm_path / "bin" / "hipconfig"
    if hipconfig.is_file():
        try:
            candidates.append(
                subprocess.check_output(
                    [str(hipconfig), "--version"],
                    text=True,
                    stderr=subprocess.DEVNULL,
                )
            )
        except (OSError, subprocess.CalledProcessError):
            pass

    for candidate in candidates:
        match = _ROCM_VERSION_RE.search(candidate)
        if match is not None:
            version = match.group(0)
            return version, f"{match.group(1)}.{match.group(2)}.*"
    return None, None


def _write_rocm_build_info(root_dir: Path) -> None:
    """Bake the build ROCm version into the package for runtime validation."""

    version, version_pattern = _detect_rocm_build_version()
    info_path = root_dir / "python" / "mori" / "_rocm_build_info.py"
    info_path.write_text(
        "# Auto-generated by setup.py at build time — do not edit.\n"
        f"ROCM_BUILD_VERSION = {version!r}\n"
        f"ROCM_VERSION_PATTERN = {version_pattern!r}\n",
        encoding="utf-8",
    )


_3RDPARTY_DIRS = ["3rdparty/spdlog", "3rdparty/msgpack-c"]


def _ensure_3rdparty(root_dir: Path, extra_dirs: list[str] | None = None) -> None:
    """Ensure 3rdparty submodule directories exist via git submodule update.

    Only the submodules in *required_dirs* are initialised.  Pass extra_dirs to
    opt-in to optional submodules.  SPDK is intentionally excluded from this
    path because tools/setup_spdk.sh does its own selective checkout.
    """
    required_dirs = _3RDPARTY_DIRS + (extra_dirs or [])
    missing = [
        d
        for d in required_dirs
        if not (root_dir / d).is_dir() or not any((root_dir / d).iterdir())
    ]
    if not missing:
        return

    for d in missing:
        (root_dir / d).mkdir(parents=True, exist_ok=True)

    try:
        subprocess.check_call(
            ["git", "config", "--global", "--add", "safe.directory", str(root_dir)],
            stdout=subprocess.DEVNULL,
            stderr=subprocess.DEVNULL,
        )
        subprocess.check_call(
            ["git", "submodule", "update", "--init", "--recursive"] + missing,
            cwd=str(root_dir),
            stdout=subprocess.DEVNULL,
        )
    except (subprocess.CalledProcessError, FileNotFoundError):
        pass

    still_missing = [d for d in required_dirs if not any((root_dir / d).iterdir())]
    if still_missing:
        raise RuntimeError(
            f"Missing 3rdparty dependencies: {still_missing}. "
            "Run 'git submodule update --init --recursive' in the source directory."
        )


def _setup_spdk(root_dir: Path) -> None:
    setup_spdk = root_dir / "tools" / "setup_spdk.sh"
    if not setup_spdk.is_file():
        raise RuntimeError(f"Missing SPDK setup script: {setup_spdk}")

    subprocess.check_call(
        [str(setup_spdk), "--jobs", str(os.cpu_count() or 1)],
        cwd=str(root_dir),
    )


class CMakeBuild(build_ext):
    def run(self) -> None:
        try:
            subprocess.check_output(["cmake", "--version"])
        except OSError as exn:
            raise RuntimeError(
                "CMake is required. Install via: pip install cmake  OR  sudo apt-get install cmake"
            ) from exn
        mpi_enabled = (
            os.environ.get("BUILD_EXAMPLES", "OFF").upper() == "ON"
            or os.environ.get("BUILD_BENCHMARK", "OFF").upper() == "ON"
            or os.environ.get("MORI_WITH_MPI", "OFF").upper() == "ON"
        )
        if mpi_enabled:
            _REQUIRED_SYSTEM_DEPS.extend(_MPI_SYSTEM_DEPS)
        _check_system_deps()
        torch_exts = [
            e for e in self.extensions if getattr(e, "_mori_torch_ext", False)
        ]
        for ext in self.extensions:
            if ext not in torch_exts:
                self.build_extension(ext)
        if torch_exts:
            try:
                self._build_with_torch(torch_exts)
            except Exception as exc:
                # AUTO means "build it if this environment can": a torch whose
                # SymmetricMemory interface has moved must not take the rest of mori
                # down with it. mori.allocator rebuilds from the shipped sources on
                # first import, so the backend is still reachable after this.
                if _torch_symm_mode() != "AUTO":
                    raise
                print(
                    f"[mori] torch SymmetricMemory backend not prebuilt ({exc}). "
                    "It will be compiled on first `import mori.allocator`; set "
                    "BUILD_TORCH_SYMM=ON to make this a build error instead."
                )

    def _build_with_torch(self, exts: list) -> None:
        """Build torch-linked extensions with torch's own build_ext, in a separate
        command instance so nothing else in this build sees torch's compiler patches.

        torch derives _GLIBCXX_USE_CXX11_ABI, its bundled pybind11 and the module's ABI
        tag from the installed torch; hand-rolling those is what we are avoiding."""
        from torch.utils.cpp_extension import BuildExtension

        cmd = BuildExtension(self.distribution)
        cmd.initialize_options()
        cmd.inplace = self.inplace
        cmd.build_lib = self.build_lib
        cmd.build_temp = self.build_temp
        cmd.force = self.force
        cmd.finalize_options()
        # after finalize_options, which would otherwise reset this to all ext_modules
        cmd.extensions = exts
        cmd.run()

    def build_extension(self, ext: Extension) -> None:
        if ext.sources and any(s.endswith((".pyx", ".cpp")) for s in ext.sources):
            if self.compiler is None:
                self.ensure_finalized()
                # Take new_compiler from the module holding the base build_ext
                # super() reaches. `distutils` and `setuptools._distutils` load
                # the same source as two separate modules, so each ends up with
                # its own CCompiler class, and build_extension (setuptools >=
                # 84) asserts isinstance against the one from its own module --
                # a compiler built from the other copy fails that check even
                # though it is otherwise fully configured.
                from setuptools._distutils.ccompiler import (
                    new_compiler as _default_new_compiler,
                )
                from setuptools._distutils.sysconfig import (
                    customize_compiler as _default_customize_compiler,
                )

                bx_mod = next(
                    (
                        sys.modules[k.__module__]
                        for k in type(self).__mro__
                        if k.__module__.endswith("distutils.command.build_ext")
                    ),
                    None,
                )
                # getattr, not attribute access: a module that matched the name
                # but re-exports neither helper must still fall back, not raise.
                new_compiler = getattr(bx_mod, "new_compiler", _default_new_compiler)
                customize_compiler = getattr(
                    bx_mod, "customize_compiler", _default_customize_compiler
                )

                try:
                    # distutils / older setuptools signature
                    self.compiler = new_compiler(
                        verbose=self.verbose,
                        dry_run=self.dry_run,
                        force=self.force,
                    )
                except TypeError:
                    # setuptools >= ~80 dropped verbose/dry_run/force kwargs
                    self.compiler = new_compiler()
                    for _attr in ("verbose", "dry_run", "force"):
                        setattr(self.compiler, _attr, getattr(self, _attr))
                customize_compiler(self.compiler)
                if self.include_dirs is not None:
                    self.compiler.set_include_dirs(self.include_dirs)
                if self.define is not None:
                    for name, value in self.define:
                        self.compiler.define_macro(name, value)
                if self.undef is not None:
                    for name in self.undef:
                        self.compiler.undefine_macro(name)
                if self.libraries is not None:
                    self.compiler.set_libraries(self.libraries)
                if self.library_dirs is not None:
                    self.compiler.set_library_dirs(self.library_dirs)
                if self.rpath is not None:
                    self.compiler.set_runtime_library_dirs(self.rpath)
                if self.link_objects is not None:
                    self.compiler.set_link_objects(self.link_objects)
            super().build_extension(ext)
            return

        build_lib = Path(self.build_lib)
        build_lib.mkdir(parents=True, exist_ok=True)

        root_dir = Path(__file__).parent

        build_umbp_spdk_enabled = _env_flag("BUILD_UMBP_SPDK", "OFF")
        build_umbp_enabled = _env_flag("BUILD_UMBP", "ON") or build_umbp_spdk_enabled

        # UMBP's distributed control plane requires gRPC + Protobuf, which are
        # system packages (not pip-installable). If they are missing, warn and
        # disable UMBP instead of letting CMake fail hard mid-configure, so the
        # rest of mori still builds and installs.
        if build_umbp_enabled:
            missing_umbp_deps = _detect_missing_umbp_deps()
            if missing_umbp_deps:
                _warn_umbp_disabled(
                    missing_umbp_deps,
                    explicit=("BUILD_UMBP" in os.environ or build_umbp_spdk_enabled),
                )
                build_umbp_enabled = False
                build_umbp_spdk_enabled = False

        _ensure_3rdparty(root_dir)
        if build_umbp_spdk_enabled:
            _setup_spdk(root_dir)

        build_dir = root_dir / os.environ.get("MORI_PYBUILD_DIR", "build")
        build_dir.mkdir(parents=True, exist_ok=True)

        cmake_cache = build_dir / "CMakeCache.txt"

        build_type = os.environ.get("CMAKE_BUILD_TYPE", "Release")
        unroll_value = os.environ.get("WARP_ACCUM_UNROLL", "1")
        build_shmem_device_wrapper = os.environ.get("BUILD_SHMEM_DEVICE_WRAPPER", "ON")
        enable_profiler = os.environ.get("ENABLE_PROFILER", "OFF")
        enable_debug_printf = os.environ.get("ENABLE_DEBUG_PRINTF", "OFF")
        enable_io_fault_injection = os.environ.get("ENABLE_IO_FAULT_INJECTION", "OFF")

        enable_standard_moe_adapt = os.environ.get("ENABLE_STANDARD_MOE_ADAPT", "OFF")
        multithread_support = os.environ.get("MORI_MULTITHREAD_SUPPORT", "OFF")
        gpu_archs = _get_gpu_archs()
        print(f"[mori] GPU architecture: {gpu_archs}")
        build_examples = os.environ.get("BUILD_EXAMPLES", "OFF")
        build_benchmark = os.environ.get("BUILD_BENCHMARK", "OFF")
        build_tests = os.environ.get("BUILD_TESTS", "OFF")
        build_umbp = "ON" if build_umbp_enabled else "OFF"
        build_umbp_spdk = "ON" if build_umbp_spdk_enabled else "OFF"
        use_redis_backend = os.environ.get("USE_REDIS_BACKEND", "OFF")
        build_xla_ffi_ops = os.environ.get("BUILD_XLA_FFI_OPS", "OFF")
        with_mpi = (
            "ON"
            if (
                build_examples.upper() == "ON"
                or build_benchmark.upper() == "ON"
                or os.environ.get("MORI_WITH_MPI", "OFF").upper() == "ON"
            )
            else "OFF"
        )
        build_ops_device = (
            "ON"
            if build_xla_ffi_ops.upper() == "ON"
            else os.environ.get("BUILD_OPS_DEVICE", "OFF")
        )
        # ON by default, matching CMakeLists.txt. See the comment there for why:
        # the flag creates no queue on its own, and OFF produces a build whose
        # SDMA puts silently move no bytes.
        #
        # Normalised, because this value is consumed twice and the two used to
        # disagree: CMake takes 1/TRUE/yes as true, while the baked-in flag was
        # `== "ON"`. `BUILD_CCO_SDMA=1` therefore built a host library with SDMA
        # and JITted a device wrapper without it -- the same silent-zeros failure
        # this default exists to prevent, from the other side.
        BUILD_CCO_SDMA = (
            "ON"
            if os.environ.get("BUILD_CCO_SDMA", "ON").strip().upper()
            in ("ON", "1", "TRUE", "YES", "Y")
            else "OFF"
        )
        if build_benchmark.upper() == "ON" and BUILD_CCO_SDMA.upper() != "ON":
            print(
                "[mori] BUILD_BENCHMARK=ON requires BUILD_CCO_SDMA=ON "
                "(cco benchmarks use ccoSdma); forcing BUILD_CCO_SDMA=ON.",
                file=sys.stderr,
            )
            BUILD_CCO_SDMA = "ON"

        cmake_args = [
            "cmake",
            "-DUSE_ROCM=ON",
            f"-DCMAKE_BUILD_TYPE={build_type}",
            f"-DWARP_ACCUM_UNROLL={unroll_value}",
            f"-DBUILD_SHMEM_DEVICE_WRAPPER={build_shmem_device_wrapper}",
            f"-DENABLE_DEBUG_PRINTF={enable_debug_printf}",
            f"-DENABLE_IO_FAULT_INJECTION={enable_io_fault_injection}",
            f"-DENABLE_STANDARD_MOE_ADAPT={enable_standard_moe_adapt}",
            f"-DGPU_TARGETS={gpu_archs}",
            f"-DENABLE_PROFILER={enable_profiler}",
            f"-DBUILD_EXAMPLES={build_examples}",
            f"-DBUILD_BENCHMARK={build_benchmark}",
            f"-DBUILD_TESTS={build_tests}",
            f"-DBUILD_UMBP={build_umbp}",
            f"-DUSE_REDIS_BACKEND={use_redis_backend}",
            f"-DUSE_SPDK={build_umbp_spdk}",
            f"-DWITH_MPI={with_mpi}",
            "-DBUILD_TORCH_BOOTSTRAP=OFF",
            f"-DBUILD_XLA_FFI_OPS={build_xla_ffi_ops}",
            f"-DBUILD_OPS_DEVICE={build_ops_device}",
            f"-DBUILD_CCO_SDMA={BUILD_CCO_SDMA}",
            f"-DMORI_MULTITHREAD_SUPPORT={multithread_support}",
            "-B",
            str(build_dir),
            "-S",
            str(root_dir),
        ]

        if shutil.which("ninja"):
            cmake_args.insert(1, "-G")
            cmake_args.insert(2, "Ninja")

        if shutil.which("ccache"):
            cmake_args.append("-DCMAKE_C_COMPILER_LAUNCHER=ccache")
            cmake_args.append("-DCMAKE_CXX_COMPILER_LAUNCHER=ccache")

        _invalidate_cmake_cache_if_changed(cmake_cache, cmake_args)
        subprocess.check_call(cmake_args)
        subprocess.check_call(
            ["cmake", "--build", ".", "-j", f"{os.cpu_count()}"], cwd=str(build_dir)
        )

        # When benchmarks are off, the shared libs are rebuilt without MPI but a
        # previous BUILD_BENCHMARK=ON run may have left benchmark executables in
        # build/benchmark/. Running those stale binaries fails with an
        # undefined MpiBootstrapNetwork symbol. Remove them so the build dir
        # stays self-consistent.
        if build_benchmark.upper() != "ON":
            bench_dir = build_dir / "benchmark"
            if bench_dir.is_dir():
                for exe in bench_dir.iterdir():
                    if exe.is_file() and os.access(exe, os.X_OK):
                        exe.unlink()

        files_to_copy = [
            (
                build_dir / "src/cco/libmori_cco.so",
                root_dir / "python/mori/libmori_cco.so",
            ),
            (
                build_dir / "src/pybind/libmori_pybinds.so",
                root_dir / "python/mori/libmori_pybinds.so",
            ),
            (
                build_dir / "src/application/libmori_application.so",
                root_dir / "python/mori/libmori_application.so",
            ),
            (
                build_dir / "src/shmem/libmori_shmem.so",
                root_dir / "python/mori/libmori_shmem.so",
            ),
            (
                build_dir / "src/ops/libmori_ops.so",
                root_dir / "python/mori/libmori_ops.so",
            ),
            (
                build_dir / "src/io/libmori_io.so",
                root_dir / "python/mori/libmori_io.so",
            ),
            (
                build_dir / "src/metrics/libmori_metrics.so",
                root_dir / "python/mori/libmori_metrics.so",
            ),
            # JIT v2. plan_api looks for these next to the mori package first, so
            # this is where they have to land; without them a wheel raises
            # "libmori_ops_v2.so not found" on the hip backend, which a dev tree
            # hides because plan_api also probes build/src/**.
            (
                build_dir / "src/jit/v2/libmori_jit.so",
                root_dir / "python/mori/libmori_jit.so",
            ),
            (
                build_dir / "src/ops/dispatch_combine_v2/libmori_ops_v2.so",
                root_dir / "python/mori/libmori_ops_v2.so",
            ),
        ]
        collective_so = build_dir / "src/collective/libmori_collective.so"
        if collective_so.exists():
            files_to_copy.append(
                (collective_so, root_dir / "python/mori/libmori_collective.so")
            )
        for src_path, dst_path in files_to_copy:
            shutil.copyfile(src_path, dst_path)

        # UMBP bindings are their own module (it links gRPC); mori.cpp loads it lazily.
        # Keyed on the flag, not the file: CMake leaves a stale one in a reused build dir.
        umbp_pybinds_src = build_dir / "src/pybind/libmori_umbp_pybinds.so"
        umbp_pybinds_dst = root_dir / "python/mori/libmori_umbp_pybinds.so"
        if build_umbp_enabled and umbp_pybinds_src.exists():
            shutil.copyfile(umbp_pybinds_src, umbp_pybinds_dst)
        elif umbp_pybinds_dst.exists():
            umbp_pybinds_dst.unlink()

        spdk_proxy_src = build_dir / "src/umbp/spdk_proxy"
        spdk_proxy_dst = root_dir / "python/mori/spdk_proxy"
        if build_umbp_spdk_enabled and spdk_proxy_src.exists():
            shutil.copyfile(spdk_proxy_src, spdk_proxy_dst)
            os.chmod(spdk_proxy_dst, 0o700)
        elif spdk_proxy_dst.exists():
            spdk_proxy_dst.unlink()

        umbp_master_src = build_dir / "src/umbp/umbp_master"
        umbp_master_dst = root_dir / "python/mori/umbp_master"
        if umbp_master_src.exists():
            shutil.copyfile(umbp_master_src, umbp_master_dst)
            os.chmod(umbp_master_dst, 0o700)
        elif umbp_master_dst.exists():
            umbp_master_dst.unlink()

        umbp_standalone_src = build_dir / "src/umbp/umbp_standalone_server"
        umbp_standalone_dst = root_dir / "python/mori/umbp_standalone_server"
        if umbp_standalone_src.exists():
            shutil.copyfile(umbp_standalone_src, umbp_standalone_dst)
            os.chmod(umbp_standalone_dst, 0o700)
        elif umbp_standalone_dst.exists():
            umbp_standalone_dst.unlink()

        # CCO C++ examples: ship the built binaries when BUILD_EXAMPLES=ON. They
        # carry an $ORIGIN/../.. rpath (set in examples/CMakeLists.txt) so they
        # resolve libmori_*.so from site-packages/mori/ once installed here.
        cco_examples_dst = root_dir / "python/mori/examples/cco"
        for _exe in ("cco_lsa_put", "cco_gda_put"):
            src = build_dir / "examples" / _exe
            dst = cco_examples_dst / _exe
            if build_examples.upper() == "ON" and src.exists():
                cco_examples_dst.mkdir(parents=True, exist_ok=True)
                shutil.copyfile(src, dst)
                os.chmod(dst, 0o755)
            elif dst.exists():
                dst.unlink()

        # CCO benchmarks: same pattern, gated on BUILD_BENCHMARK=ON.
        cco_bench_dst = root_dir / "python/mori/benchmarks/cco"
        for _exe in (
            "cco_p2p_put_bw",
            "cco_p2p_put_latency",
            "cco_p2p_get_bw",
            "cco_p2p_get_latency",
        ):
            src = build_dir / "benchmark" / _exe
            dst = cco_bench_dst / _exe
            if build_benchmark.upper() == "ON" and src.exists():
                cco_bench_dst.mkdir(parents=True, exist_ok=True)
                shutil.copyfile(src, dst)
                os.chmod(dst, 0o755)
            elif dst.exists():
                dst.unlink()

        _copy_jit_sources(root_dir)

        # Persist the SDMA compile flag into the package so the runtime JIT
        # (mori.cco.device.bitcode) defaults to the SAME value the host lib was
        # built with — without this, a `BUILD_CCO_SDMA=ON pip install .` build
        # would still JIT an OFF device wrapper unless the user re-set the env at
        # runtime. The env var still overrides this baked-in default.
        _write_cco_build_flags(root_dir, sdma=BUILD_CCO_SDMA.upper() == "ON")
        _write_rocm_build_info(root_dir)

        if os.environ.get("MORI_SKIP_PRECOMPILE", "").lower() not in (
            "1",
            "true",
            "on",
        ):
            _try_precompile(root_dir)


def _try_precompile(root_dir: Path) -> None:
    """Precompile JIT kernels in the background if a GPU is detected.

    Launches a detached subprocess that compiles all .hsaco kernels and shmem
    bitcode into ~/.mori/jit/. The subprocess is fire-and-forget — pip install
    returns immediately without waiting.

    If the user starts using kernels before precompilation finishes, the JIT
    framework handles the race safely via FileBaton file locks: the user process
    either waits for the background compile to finish, or compiles the kernel
    itself (the background process will skip already-compiled kernels).
    """
    if _detect_local_gpu_arch() is None:
        print("[mori] No GPU detected — skipping kernel precompilation")
        return
    rocm_path = os.environ.get("ROCM_PATH", "/opt/rocm")
    hipcc = os.path.join(rocm_path, "bin", "hipcc")
    if not os.path.isfile(hipcc):
        print(f"[mori] hipcc not found at {hipcc} — skipping kernel precompilation")
        return
    try:
        target_python = os.environ.get(
            "MORI_PYTHON",
            shutil.which("python3") or shutil.which("python") or sys.executable,
        )
        env = os.environ.copy()
        env["MORI_PRECOMPILE"] = "1"
        env.pop("PYTHONPATH", None)
        subprocess.Popen(
            [target_python, "-c", "import time; time.sleep(3); import mori"],
            env=env,
            stdout=subprocess.DEVNULL,
            stderr=subprocess.DEVNULL,
        )
        print("[mori] Kernel precompilation started in background")
    except Exception as e:
        print(f"[mori] Precompilation skipped: {e}")


_BUNDLED_TOOL_FILES = ("env_check.sh", "env_setup.sh", "diagnose_env.sh", "mori.conf")


def _sync_bundled_scripts() -> None:
    """Copy bundled tools into ``python/mori/tools/`` so they ship in the wheel.

    Keeps a single source of truth (``tools/``) while still letting the
    installed package expose them via the ``mori`` console script.
    """
    here = Path(__file__).resolve().parent
    src_dir = here / "tools"
    dst_dir = here / "python" / "mori" / "tools"
    dst_dir.mkdir(parents=True, exist_ok=True)
    for name in _BUNDLED_TOOL_FILES:
        src = src_dir / name
        if not src.is_file():
            continue
        dst = dst_dir / name
        try:
            if not dst.is_file() or dst.read_bytes() != src.read_bytes():
                shutil.copy2(src, dst)
            if src.suffix == ".sh":
                os.chmod(dst, os.stat(dst).st_mode | 0o111)
        except OSError as exc:
            print(f"[mori] WARN: failed to bundle {name}: {exc}")


_sync_bundled_scripts()


class CustomBuild(_build):
    def run(self) -> None:
        _sync_bundled_scripts()
        self.run_command("build_ext")
        super().run()


_root_dir = Path(__file__).parent


def _cco_extension() -> list:
    """Build the mori.cco.cco Cython C++ extension if Cython is available."""
    if not _HAVE_CYTHON:
        print(
            "[mori] WARNING: Cython not found — skipping mori.cco.cco extension. "
            "Install Cython to enable: pip install cython",
            file=sys.stderr,
        )
        return []
    include_dirs = [str(_root_dir / "include")]
    library_dirs = [str(_root_dir / "python/mori")]
    ext = Extension(
        "mori.cco.cco",
        sources=["python/mori/cco/cco.pyx"],
        language="c++",
        include_dirs=include_dirs,
        library_dirs=library_dirs,
        libraries=["mori_cco"],
        runtime_library_dirs=["$ORIGIN/.."],
        extra_compile_args=["-std=c++17"],
    )
    return _cythonize(
        [ext],
        compiler_directives={"language_level": "3"},
    )


def _torch_symm_extension():
    """SymmetricMemory backend. Built by torch's cpp_extension, not CMake: it is the
    only target that links libtorch, and torch's build_ext is what keeps the ABI flag,
    pybind11 copy and module suffix consistent with the installed torch."""
    mode = _torch_symm_mode()
    if mode == "OFF":
        return []
    # rendezvous() is a cco window now, so the extension links libmori_cco. There is
    # deliberately no BUILD_CCO check here: that env var is not forwarded to cmake
    # (only BUILD_CCO_SDMA is), so reading it would gate on a flag that does not
    # control whether libmori_cco is built. A build that really has no libmori_cco
    # -- CMAKE_ARGS="-DBUILD_CCO=OFF" -- fails at link, which AUTO already reports
    # and skips, and which BUILD_TORCH_SYMM=ON already turns into a hard error.
    if _TorchCppExtension is None:
        if mode == "ON":
            raise RuntimeError(
                "BUILD_TORCH_SYMM=ON but torch is not importable at build time"
            )
        return []
    rocm = os.environ.get("ROCM_PATH", "/opt/rocm")
    # Absolute: torch's ninja compiler writes its build file into build_temp and runs
    # ninja from there, so a relative -I would resolve against the wrong directory.
    ext = _TorchCppExtension(
        name="mori.mori_torch_symm",
        sources=["src/allocator/symm_backend.cpp"],
        include_dirs=[str(_root_dir.resolve() / "include"), f"{rocm}/include"],
        # libmori_cco is the peer-exchange half of the backend; it lands next to this
        # module in the package, hence the $ORIGIN runpath.
        library_dirs=[f"{rocm}/lib", str(_root_dir.resolve() / "python" / "mori")],
        libraries=["amdhip64", "c10_hip", "torch_hip", "mori_cco"],
        runtime_library_dirs=["$ORIGIN"],
        # No -std here: torch's BuildExtension adds the standard its headers need.
    )
    ext._mori_torch_ext = True
    return [ext]


extensions = (
    [
        Extension(
            "mori",
            sources=[],
            # extra_compile_args=['-ggdb', '-O0'],
            # extra_link_args=['-g'],
        ),
    ]
    + _cco_extension()
    + _torch_symm_extension()
)

mori_package_data = [
    "libmori_cco.so",
    "libmori_pybinds.so",
    "libmori_shmem.so",
    "libmori_ops.so",
    "libmori_io.so",
    "libmori_application.so",
    "libmori_metrics.so",
    "libmori_collective.so",  # optional: only present when BUILD_COLLECTIVE=ON
    "libmori_jit.so",  # JIT v2 C ABI + plan registry
    "libmori_ops_v2.so",  # v2 EP kernels' host side; registers into libmori_jit
    "libmori_umbp_pybinds.so",  # optional: only present when BUILD_UMBP=ON
    "umbp_master",
    "umbp_standalone_server",
    "_jit-sources/include/**/*.hpp",
    "_jit-sources/include/**/*.h",
    "_jit-sources/include/**/*.cuh",
    "_jit-sources/src/**/*.hip",
    "_jit-sources/src/**/*.hpp",
    "_jit-sources/src/**/*.cpp",
    "_jit-sources/src/**/*.h",
    "_jit-sources/3rdparty/**/*.h",
    "_jit-sources/3rdparty/**/*.hpp",
    "_jit-sources/tools/**/*.py",
    "ops/tuning_configs/*.json",
    "tools/*.sh",
    "tools/mori.conf",
    "examples/cco/*",  # CCO C++ example binaries (only present when BUILD_EXAMPLES=ON)
    "benchmarks/cco/*",  # CCO benchmark binaries (only present when BUILD_BENCHMARK=ON)
]
if _env_flag("BUILD_UMBP_SPDK", "OFF"):
    mori_package_data.append("spdk_proxy")

setup(
    packages=find_packages(where="python"),
    package_dir={"": "python"},
    package_data={
        "mori": mori_package_data,
        "mori.cco": ["*.pxd"],
        "mori.cco.device": ["*.bc"],
        "mori.ir": ["*.bc"],
        "mori.tools": ["*.sh"],
    },
    exclude_package_data={
        "mori": ["*.a"],
    },
    cmdclass={
        "build_ext": CMakeBuild,
        "build": CustomBuild,
    },
    ext_modules=extensions,
    include_package_data=True,
)