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
76 lines
2.3 KiB
Python
76 lines
2.3 KiB
Python
from pathlib import Path
|
|
|
|
import pytest
|
|
from flashinfer_jit_cache_installer import (
|
|
expected_jit_cache_version,
|
|
install_url_if_needed,
|
|
jit_cache_wheel_url,
|
|
read_exact_pin,
|
|
)
|
|
|
|
|
|
def test_read_exact_pin_ignores_other_requirements(tmp_path: Path):
|
|
requirements = tmp_path / "cuda.txt"
|
|
requirements.write_text(
|
|
"-r common.txt\n"
|
|
"torch==2.11.0\n"
|
|
"flashinfer-python==0.6.13\n"
|
|
"flashinfer-cubin==0.6.13\n"
|
|
)
|
|
|
|
assert read_exact_pin(requirements, "flashinfer-python") == "0.6.13"
|
|
|
|
|
|
def test_read_exact_pin_requires_exact_pin(tmp_path: Path):
|
|
requirements = tmp_path / "cuda.txt"
|
|
requirements.write_text("flashinfer-python>=0.6\n")
|
|
|
|
with pytest.raises(ValueError, match="flashinfer-python exact pin not found"):
|
|
read_exact_pin(requirements, "flashinfer-python")
|
|
|
|
|
|
def test_jit_cache_url_tracks_flashinfer_and_cuda_versions():
|
|
assert expected_jit_cache_version("0.6.13", "130") == "0.6.13+cu130"
|
|
assert jit_cache_wheel_url("0.6.13", "130") == (
|
|
"https://github.com/flashinfer-ai/flashinfer/releases/download/"
|
|
"v0.6.13/flashinfer_jit_cache-0.6.13+cu130-cp39-abi3-"
|
|
"manylinux_2_28_aarch64.whl"
|
|
)
|
|
|
|
|
|
def test_install_url_if_needed_skips_matching_version(tmp_path: Path):
|
|
requirements = tmp_path / "cuda.txt"
|
|
requirements.write_text("flashinfer-python==0.6.13\n")
|
|
|
|
url, expected, installed = install_url_if_needed(
|
|
requirements,
|
|
"130",
|
|
installed_version="0.6.13+cu130",
|
|
)
|
|
|
|
assert url is None
|
|
assert expected == "0.6.13+cu130"
|
|
assert installed == "0.6.13+cu130"
|
|
|
|
|
|
def test_install_url_if_needed_reinstalls_missing_or_stale_version(tmp_path: Path):
|
|
requirements = tmp_path / "cuda.txt"
|
|
requirements.write_text("flashinfer-python==0.6.13\n")
|
|
|
|
missing_url, _, missing_installed = install_url_if_needed(
|
|
requirements,
|
|
"130",
|
|
installed_version=None,
|
|
)
|
|
stale_url, _, stale_installed = install_url_if_needed(
|
|
requirements,
|
|
"130",
|
|
installed_version="0.6.11.post3+cu130",
|
|
)
|
|
|
|
expected_url = jit_cache_wheel_url("0.6.13", "130")
|
|
assert missing_url == expected_url
|
|
assert missing_installed is None
|
|
assert stale_url == expected_url
|
|
assert stale_installed == "0.6.11.post3+cu130"
|