wehub-resource-sync
59a0a3844c
PR Test AMD / cancel-on-close (push) Has been skipped
PR Test NVIDIA ARM / scan (push) Has been skipped
PR Test NVIDIA / cancel-on-close (push) Has been skipped
PR Test AMD / scan (push) Has been skipped
PR Test NVIDIA ARM / cancel-on-close (push) Has been skipped
PR Test NVIDIA / scan (push) Has been skipped
Release Docker Images / build (cu129-torch-2.11.0) (push) Has been skipped
Release Docker Images / build (cu130-torch-2.11.0) (push) Has been skipped
Release PyPI / publish (push) Has been skipped
Scheduler Python Test / test (push) Successful in 27m19s
Docs / build (push) Successful in 28m8s
Scheduler C++ Test / test (push) Successful in 28m19s
Scheduler C++ Test / test-flat (push) Successful in 28m18s
Docs / deploy (push) Has been cancelled
PR Test AMD / finish (push) Has been cancelled
PR Test NVIDIA / finish (push) Has been cancelled
PR Test NVIDIA ARM / finish (push) Has been cancelled
PR Test NVIDIA ARM / ${{ matrix.name }} (${{ matrix.runner }}) (push) Has been cancelled
PR Test AMD / ${{ matrix.name }} (${{ matrix.runner }}) (push) Has been cancelled
PR Test NVIDIA / ${{ matrix.name }} (${{ matrix.runner }}) (push) Has been cancelled
66 lines
2.3 KiB
Python
66 lines
2.3 KiB
Python
from __future__ import annotations
|
|
|
|
import runpy
|
|
from pathlib import Path
|
|
|
|
import setuptools
|
|
|
|
SETUP_PY = Path(__file__).parents[1] / "python" / "setup.py"
|
|
|
|
|
|
def test_cuda_include_dirs_prefer_complete_toolkit_headers(
|
|
tmp_path, monkeypatch
|
|
) -> None:
|
|
monkeypatch.setenv("TOKENSPEED_KERNEL_BACKEND", "cuda")
|
|
monkeypatch.setattr(setuptools, "setup", lambda **_kwargs: None)
|
|
setup_namespace = runpy.run_path(str(SETUP_PY))
|
|
|
|
cuda_root = tmp_path / "cuda"
|
|
cuda_include = cuda_root / "include"
|
|
cuda_include.mkdir(parents=True)
|
|
(cuda_include / "cuda_runtime.h").touch()
|
|
(cuda_include / "cublas_v2.h").touch()
|
|
|
|
site_packages = tmp_path / "site-packages"
|
|
wheel_include = site_packages / "nvidia" / "cu13" / "include"
|
|
wheel_include.mkdir(parents=True)
|
|
(wheel_include / "cuda_runtime.h").touch()
|
|
(wheel_include / "cublas_v2.h").touch()
|
|
|
|
builder = setup_namespace["CudaKernelBuilder"]([], verbose=False)
|
|
monkeypatch.setattr(builder, "_cuda_toolkit_roots", lambda: iter([cuda_root]))
|
|
monkeypatch.setattr(builder, "_site_paths", lambda: iter([site_packages]))
|
|
|
|
include_dirs = builder._resolve_include_dirs()
|
|
|
|
assert str(cuda_include) in include_dirs
|
|
assert str(wheel_include) not in include_dirs
|
|
|
|
|
|
def test_cuda_include_dirs_fall_back_from_partial_toolkit(
|
|
tmp_path, monkeypatch
|
|
) -> None:
|
|
monkeypatch.setenv("TOKENSPEED_KERNEL_BACKEND", "cuda")
|
|
monkeypatch.setattr(setuptools, "setup", lambda **_kwargs: None)
|
|
setup_namespace = runpy.run_path(str(SETUP_PY))
|
|
|
|
cuda_root = tmp_path / "cuda"
|
|
cuda_include = cuda_root / "include"
|
|
cuda_include.mkdir(parents=True)
|
|
(cuda_include / "cuda_runtime.h").touch()
|
|
|
|
site_packages = tmp_path / "site-packages"
|
|
wheel_include = site_packages / "nvidia" / "cu13" / "include"
|
|
wheel_include.mkdir(parents=True)
|
|
(wheel_include / "cuda_runtime.h").touch()
|
|
(wheel_include / "cublas_v2.h").touch()
|
|
|
|
builder = setup_namespace["CudaKernelBuilder"]([], verbose=False)
|
|
monkeypatch.setattr(builder, "_cuda_toolkit_roots", lambda: iter([cuda_root]))
|
|
monkeypatch.setattr(builder, "_site_paths", lambda: iter([site_packages]))
|
|
|
|
include_dirs = builder._resolve_include_dirs()
|
|
|
|
assert str(cuda_include) not in include_dirs
|
|
assert str(wheel_include) in include_dirs
|