Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
61 changes: 38 additions & 23 deletions cuda_core/build_hooks.py
Original file line number Diff line number Diff line change
Expand Up @@ -11,7 +11,6 @@
import glob
import os
import re
import subprocess

from Cython.Build import cythonize
from setuptools import Extension
Expand All @@ -26,32 +25,48 @@


@functools.cache
def _get_proper_cuda_bindings_major_version() -> str:
# for local development (with/without build isolation)
try:
import cuda.bindings
def _get_cuda_major_version() -> str:
"""Determine the CUDA major version for building cuda.core.

return cuda.bindings.__version__.split(".")[0]
except ImportError:
pass
This version is used for two purposes:
1. Determining which cuda-bindings version to install as a build dependency
2. Setting CUDA_CORE_BUILD_MAJOR for Cython compile-time conditionals

# for custom overwrite, e.g. in CI
The version is derived from (in order of priority):
1. CUDA_CORE_BUILD_MAJOR environment variable (explicit override, e.g. in CI)
2. CUDA_VERSION macro in cuda.h from CUDA_PATH or CUDA_HOME

Since CUDA_PATH or CUDA_HOME is required for the build (to provide include
directories), the cuda.h header should always be available.
"""
# Explicit override, e.g. in CI.
cuda_major = os.environ.get("CUDA_CORE_BUILD_MAJOR")
if cuda_major is not None:
return cuda_major

# also for local development
try:
out = subprocess.run("nvidia-smi", env=os.environ, capture_output=True, check=True) # noqa: S603, S607
m = re.search(r"CUDA Version:\s*([\d\.]+)", out.stdout.decode())
if m:
return m.group(1).split(".")[0]
except (FileNotFoundError, subprocess.CalledProcessError):
# the build machine has no driver installed
pass

# default fallback
return "13"
# Derive from the CUDA headers (the authoritative source for what we compile against).
cuda_path = os.environ.get("CUDA_PATH", os.environ.get("CUDA_HOME", None))
if cuda_path:
for root in cuda_path.split(os.pathsep):
cuda_h = os.path.join(root, "include", "cuda.h")
try:
with open(cuda_h, encoding="utf-8") as f:
for line in f:
m = re.match(r"^#\s*define\s+CUDA_VERSION\s+(\d+)\s*$", line)
if m:
v = int(m.group(1))
# CUDA_VERSION is e.g. 12020 for 12.2.
return str(v // 1000)
except OSError:
continue

# CUDA_PATH or CUDA_HOME is required for the build, so we should not reach here
# in normal circumstances. Raise an error to make the issue clear.
raise RuntimeError(
"Cannot determine CUDA major version. "
"Set CUDA_CORE_BUILD_MAJOR environment variable, or ensure CUDA_PATH or CUDA_HOME "
"points to a valid CUDA installation with include/cuda.h."
)


# used later by setup()
Expand Down Expand Up @@ -105,7 +120,7 @@ def get_cuda_paths():
)

nthreads = int(os.environ.get("CUDA_PYTHON_PARALLEL_LEVEL", os.cpu_count() // 2))
compile_time_env = {"CUDA_CORE_BUILD_MAJOR": int(_get_proper_cuda_bindings_major_version())}
compile_time_env = {"CUDA_CORE_BUILD_MAJOR": int(_get_cuda_major_version())}
compiler_directives = {"embedsignature": True, "warn.deprecated.IF": False, "freethreading_compatible": True}
if COMPILE_FOR_COVERAGE:
compiler_directives["linetrace"] = True
Expand All @@ -132,7 +147,7 @@ def build_wheel(wheel_directory, config_settings=None, metadata_directory=None):


def _get_cuda_bindings_require():
cuda_major = _get_proper_cuda_bindings_major_version()
cuda_major = _get_cuda_major_version()
return [f"cuda-bindings=={cuda_major}.*"]


Expand Down
128 changes: 128 additions & 0 deletions cuda_core/tests/test_build_hooks.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,128 @@
# SPDX-FileCopyrightText: Copyright (c) 2024-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0

"""Tests for build_hooks.py build infrastructure.

These tests verify the CUDA version detection logic used during builds,
particularly the _get_cuda_major_version() function which derives the
CUDA major version from headers.

Note: These tests do NOT require cuda.core to be built/installed since they
test build-time infrastructure. Run with --noconftest to avoid loading
conftest.py which imports cuda.core modules:

pytest tests/test_build_hooks.py -v --noconftest

These tests require Cython to be installed (build_hooks.py imports it).
"""

import importlib.util
import os
import tempfile
from pathlib import Path
from unittest import mock

import pytest

# build_hooks.py imports Cython at the top level, so skip if not available
pytest.importorskip("Cython")


def _load_build_hooks():
"""Load build_hooks module from source without permanently modifying sys.path.

build_hooks.py is a PEP 517 build backend, not an installed module.
We use importlib to load it directly from source to avoid polluting
sys.path with the cuda_core/ directory (which contains cuda/core/ source
that could shadow the installed package).
"""
build_hooks_path = Path(__file__).parent.parent / "build_hooks.py"
spec = importlib.util.spec_from_file_location("build_hooks", build_hooks_path)
module = importlib.util.module_from_spec(spec)
spec.loader.exec_module(module)
return module


# Load the module once at import time
build_hooks = _load_build_hooks()


def _check_version_detection(
cuda_version, expected_major, *, use_cuda_path=True, use_cuda_home=False, cuda_core_build_major=None
):
"""Test version detection with a mock cuda.h.

Args:
cuda_version: CUDA_VERSION to write in mock cuda.h (e.g., 12080)
expected_major: Expected return value (e.g., "12")
use_cuda_path: If True, set CUDA_PATH to the mock headers directory
use_cuda_home: If True, set CUDA_HOME to the mock headers directory
cuda_core_build_major: If set, override with this CUDA_CORE_BUILD_MAJOR env var
"""
with tempfile.TemporaryDirectory() as tmpdir:
include_dir = Path(tmpdir) / "include"
include_dir.mkdir()
cuda_h = include_dir / "cuda.h"
cuda_h.write_text(f"#define CUDA_VERSION {cuda_version}\n")

build_hooks._get_cuda_major_version.cache_clear()

mock_env = {
k: v
for k, v in {
"CUDA_CORE_BUILD_MAJOR": cuda_core_build_major,
"CUDA_PATH": tmpdir if use_cuda_path else None,
"CUDA_HOME": tmpdir if use_cuda_home else None,
}.items()
if v is not None
}

with mock.patch.dict(os.environ, mock_env, clear=True):
result = build_hooks._get_cuda_major_version()
assert result == expected_major


class TestGetCudaMajorVersion:
"""Tests for _get_cuda_major_version()."""

@pytest.mark.parametrize("version", ["11", "12", "13", "14"])
def test_env_var_override(self, version):
"""CUDA_CORE_BUILD_MAJOR env var override works with various versions."""
build_hooks._get_cuda_major_version.cache_clear()
with mock.patch.dict(os.environ, {"CUDA_CORE_BUILD_MAJOR": version}, clear=False):
result = build_hooks._get_cuda_major_version()
assert result == version

@pytest.mark.parametrize(
("cuda_version", "expected_major"),
[
(11000, "11"), # CUDA 11.0
(11080, "11"), # CUDA 11.8
(12000, "12"), # CUDA 12.0
(12020, "12"), # CUDA 12.2
(12080, "12"), # CUDA 12.8
(13000, "13"), # CUDA 13.0
(13010, "13"), # CUDA 13.1
],
ids=["11.0", "11.8", "12.0", "12.2", "12.8", "13.0", "13.1"],
)
def test_cuda_headers_parsing(self, cuda_version, expected_major):
"""CUDA_VERSION is correctly parsed from cuda.h headers."""
_check_version_detection(cuda_version, expected_major)

def test_cuda_home_fallback(self):
"""CUDA_HOME is used if CUDA_PATH is not set."""
_check_version_detection(12050, "12", use_cuda_path=False, use_cuda_home=True)

def test_env_var_takes_priority_over_headers(self):
"""Env var override takes priority even when headers exist."""
_check_version_detection(12080, "11", cuda_core_build_major="11")

def test_missing_cuda_path_raises_error(self):
"""RuntimeError is raised when CUDA_PATH/CUDA_HOME not set and no env var override."""
build_hooks._get_cuda_major_version.cache_clear()
with (
mock.patch.dict(os.environ, {}, clear=True),
pytest.raises(RuntimeError, match="Cannot determine CUDA major version"),
):
build_hooks._get_cuda_major_version()
Loading