setup.py

"""ffpa-attn build script.

Most package metadata (name, version, dependencies, extras, URLs, etc.) lives
in ``pyproject.toml``. This ``setup.py`` exists only to drive the optional
``CUDAExtension`` build for the ``ffpa_attn._C`` C++/CUDA module via the
PyTorch build helpers.

Behavior:
- Default: build the CUDA extension via ``torch.utils.cpp_extension``.
- ``FFPA_SKIP_CUDA_EXT=1``: skip CUDA extension entirely (used by the
  ReadTheDocs build and by the docs CI runner, which have no CUDA toolchain
  but still need to import ``ffpa_attn`` so ``mkdocstrings`` can render API
  documentation from docstrings).
"""

import os
import subprocess
import sys
import warnings
from pathlib import Path

from setuptools import setup

warnings.filterwarnings("ignore")

# Ensure the project root (containing ``env.py``) is on ``sys.path`` so that
# the build backend (``setuptools.build_meta``) can import ``env`` even when
# pip invokes setup.py from a different working directory.
_ROOT = Path(__file__).resolve().parent
if str(_ROOT) not in sys.path:
  sys.path.insert(0, str(_ROOT))

_VERSION_FILE = _ROOT / "src" / "ffpa_attn" / "_version.py"


def _resolve_version() -> str:
  try:
    from setuptools_scm import get_version as _get_scm_version

    return _get_scm_version(
      root=str(_ROOT),
      version_scheme="python-simplified-semver",
      local_scheme="no-local-version",
      write_to=str(_VERSION_FILE),
    )
  except Exception:
    pass

  try:
    describe = subprocess.check_output(
      ["git", "describe", "--tags", "--long", "--dirty", "--match", "v*"],
      cwd=_ROOT,
      text=True,
      stderr=subprocess.DEVNULL,
    ).strip()
    tag, distance, _sha, *dirty = describe.split("-")
    base = tag.removeprefix("v")
    if distance == "0" and not dirty:
      return base
    return f"{base}.dev{distance}"
  except Exception:
    return "0.0.0"


def _write_version_file(version: str) -> None:
  _VERSION_FILE.write_text(
    '# file generated by setup.py\n'
    f'version = "{version}"\n'
    '__version__ = version\n',
    encoding="utf-8",
  )


_RESOLVED_VERSION = _resolve_version()
if not _VERSION_FILE.exists():
  _write_version_file(_RESOLVED_VERSION)


def _env_flag(name: str) -> bool:
  return os.getenv(name, "0").strip().lower() in {"1", "true", "yes", "on"}


SKIP_CUDA_EXT = _env_flag("FFPA_SKIP_CUDA_EXT")

ext_modules = []
cmdclass = {}

# Force a PyPI-acceptable platform tag for the produced wheel.
#
# setuptools' default ``bdist_wheel`` emits ``linux_x86_64`` for CUDA
# extensions, which PyPI rejects (only ``manylinux*`` / ``musllinux*`` tags
# are allowed). The actual glibc symbol floor of this extension is
# determined by the build host; ``manylinux_2_34_x86_64`` matches the
# ``auditwheel show`` output on the canonical build environment (Ubuntu
# 24.04 + CUDA 13) and is also the platform PyTorch publishes against,
# so downstream ``pip install`` resolves cleanly. Override per-build via
# ``FFPA_PLAT_TAG=manylinux_2_28_x86_64 python -m build --wheel`` etc.
_PLAT_TAG = os.getenv("FFPA_PLAT_TAG", "manylinux_2_34_x86_64").strip()
if _PLAT_TAG:
  try:
    from wheel.bdist_wheel import bdist_wheel as _bdist_wheel

    class _ManylinuxBdistWheel(_bdist_wheel):

      def finalize_options(self):
        super().finalize_options()
        # Mark the wheel as platform-specific (not pure-python) and pin the
        # platform tag so PyPI accepts the upload.
        self.root_is_pure = False
        self.plat_name_supplied = True
        self.plat_name = _PLAT_TAG

    cmdclass["bdist_wheel"] = _ManylinuxBdistWheel
  except ImportError:
    # ``wheel`` is a build-time dep declared in pyproject.toml; this guard
    # only triggers in odd environments where it is genuinely missing.
    pass

if not SKIP_CUDA_EXT:
  from env import ENV
  from torch.utils.cpp_extension import BuildExtension, CUDAExtension

  ENV.list_ffpa_env()

  cc_flag = []
  for _sm in ENV.get_build_arch_list():
    cc_flag.append("-gencode")
    cc_flag.append(f"arch=compute_{_sm},code=sm_{_sm}")

  ext_modules.append(
    CUDAExtension(
      # Package-internal C extension module; imported as ``ffpa_attn._C``.
      name="ffpa_attn._C",
      sources=[
        # Convert to repo-relative paths; setuptools rejects absolute paths
        # in ``sources`` for editable installs (``pip install -e .``).
        os.path.relpath(s, _ROOT) for s in ENV.get_build_sources(build_pkg=True)
      ],
      extra_compile_args={
        "cxx": [flag for flag in ENV.extra_gcc_flags() if flag.strip()],
        "nvcc": [flag for flag in (ENV.get_build_cuda_cflags(build_pkg=True) + cc_flag) if flag.strip()],
      },
      include_dirs=[
        Path(ENV.project_dir()) / "csrc" / "cuffpa",
      ],
    )
  )
  cmdclass["build_ext"] = BuildExtension

setup(
  version=_RESOLVED_VERSION,
  ext_modules=ext_modules,
  cmdclass=cmdclass,
)