setup.py

"""
novatorch build script: drives CMake to compile the CUDA extension.

    pip install -e .                      # editable install
    pip install .                         # regular install
    python setup.py build_ext --inplace   # just build the extension

Requires the CUDA toolkit (nvcc) and CMake >= 3.20 on PATH.
"""
import os
import shutil
import subprocess
import sys
from pathlib import Path

from setuptools import Extension, setup
from setuptools.command.build_ext import build_ext

HERE = Path(__file__).parent.absolute()


def detect_cuda_arch():
    """Compute capability of the local GPU, e.g. '89'. Falls back to a fat build."""
    env = os.environ.get("NOVATORCH_CUDA_ARCH")
    if env:
        return env
    try:
        out = subprocess.check_output(
            ["nvidia-smi", "--query-gpu=compute_cap", "--format=csv,noheader"],
            stderr=subprocess.DEVNULL, text=True).strip()
        caps = {line.strip().replace(".", "") for line in out.splitlines() if line.strip()}
        if caps:
            return ";".join(sorted(caps))
    except Exception:
        pass
    return "75;80;86;89"


class CMakeBuild(build_ext):
    def run(self):
        try:
            subprocess.check_call(["cmake", "--version"],
                                  stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
        except OSError as err:
            raise RuntimeError("CMake is required to build novatorch") from err
        for ext in self.extensions:
            self.build_extension(ext)

    def build_extension(self, ext):
        build_dir = HERE / "build"
        build_dir.mkdir(parents=True, exist_ok=True)

        cmake_args = [
            f"-DPython3_EXECUTABLE={sys.executable}",
            f"-DCMAKE_CUDA_ARCHITECTURES={detect_cuda_arch()}",
            "-DCMAKE_BUILD_TYPE=Release",
        ]
        try:
            import pybind11
            cmake_args.append(f"-Dpybind11_DIR={pybind11.get_cmake_dir()}")
        except ImportError:
            pass   # CMake falls back to FetchContent

        if sys.platform == "win32":
            cmake_args.append("-DCMAKE_GENERATOR_PLATFORM=x64")
            build_args = ["--config", "Release", "--parallel"]
        else:
            build_args = ["--parallel", str(os.cpu_count() or 4)]

        subprocess.check_call(["cmake", "-S", str(HERE), "-B", str(build_dir)] + cmake_args)
        subprocess.check_call(["cmake", "--build", str(build_dir),
                               "--target", "_novatorch_C"] + build_args)

        # CMake writes the module into ./novatorch; copy it where setuptools expects it
        built = sorted((HERE / "novatorch").glob("_novatorch_C*"))
        if not built:
            raise RuntimeError("the extension was not produced by the CMake build")
        dest = Path(self.get_ext_fullpath(ext.name)).parent
        dest.mkdir(parents=True, exist_ok=True)
        for path in built:
            if path.resolve() != (dest / path.name).resolve():
                shutil.copy2(path, dest / path.name)


setup(
    name="novatorch",
    version="1.0.0",
    description="A CUDA deep learning framework with autograd, written from scratch",
    long_description=(HERE / "README.md").read_text(encoding="utf-8", errors="ignore"),
    long_description_content_type="text/markdown",
    license="MIT",
    packages=["novatorch", "novatorch.nn", "novatorch.optim", "novatorch.data"],
    package_data={"novatorch": ["*.pyd", "*.so", "*.dll"]},
    ext_modules=[Extension("novatorch._novatorch_C", sources=[])],
    cmdclass={"build_ext": CMakeBuild},
    python_requires=">=3.9",
    install_requires=["numpy>=1.21"],
    extras_require={
        "dashboard": ["fastapi>=0.100", "uvicorn>=0.23", "websockets>=11", "pynvml>=11"],
        "dev": ["pytest>=7", "pybind11>=2.12", "cmake>=3.20"],
    },
    classifiers=[
        "Development Status :: 4 - Beta",
        "Intended Audience :: Science/Research",
        "License :: OSI Approved :: MIT License",
        "Topic :: Scientific/Engineering :: Artificial Intelligence",
        "Programming Language :: Python :: 3",
        "Programming Language :: C++",
        "Environment :: GPU :: NVIDIA CUDA",
    ],
)