pypi_variants.py

"""Create Package variants for PyPI distribution."""

import argparse
import os

from test_utils import PY_PACKAGE

IN_PATH = os.path.join(PY_PACKAGE, "pyproject.toml.in")
OUT_PATH = os.path.join(PY_PACKAGE, "pyproject.toml")

CHOICES = ["default", "cpu", "manylinux2014"]


NCCL_WHL = """    \"nvidia-nccl-cu12 ; platform_system == 'Linux' and platform_machine != 'aarch64'\","""

NAME = "{{ name }}"
NCCL = "{{ nccl }}"


def copyfile(src: str, dst: str) -> None:
    with open(src, "rb") as fd:
        content = fd.read()
    with open(dst, "wb") as fd:
        fd.write(content)


def make_pyproject(variant: str) -> None:
    assert variant in CHOICES

    with open(IN_PATH) as fd:
        pyproject = fd.read()

    readme_dft = os.path.join(PY_PACKAGE, "README.dft.rst")
    readme_cpu = os.path.join(PY_PACKAGE, "README.cpu.rst")
    readme = os.path.join(PY_PACKAGE, "README.rst")
    if variant == "cpu":
        pyproject = pyproject.replace(NAME, "xgboost-cpu").replace(NCCL, "")
        copyfile(readme_cpu, readme)
    elif variant == "manylinux2014":
        pyproject = pyproject.replace(NAME, "xgboost").replace(NCCL, "")
        copyfile(readme_dft, readme)
    else:
        pyproject = pyproject.replace(NAME, "xgboost").replace(NCCL, NCCL_WHL)
        copyfile(readme_dft, readme)
    pyproject = (
        f"# Generated by `{os.path.basename(__file__)}`, don't edit.\n" + pyproject
    )

    with open(OUT_PATH, "w") as fd:
        fd.write(pyproject)


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--variant",
        type=str,
        choices=CHOICES,
        default="default",
    )
    args = parser.parse_args()
    make_pyproject(args.variant)