From d09357f5411783a5a18fdfb3aa5d19d49565cc7b Mon Sep 17 00:00:00 2001 From: gouyt13clear Date: Thu, 29 Oct 2026 13:55:56 +0800 Subject: [PATCH] fix(ci): select compatible Windows runtimes and clean platform checks --- .github/scripts/ci_changes.py | 5 + .github/scripts/repair_windows_wheel.py | 152 ++++++++++++++++++ .github/scripts/test_ci_changes.py | 5 + .github/scripts/test_ci_workflow_contracts.py | 5 +- .github/scripts/test_repair_windows_wheel.py | 137 ++++++++++++++++ pyproject.toml | 2 +- tests/CMakeLists.txt | 4 +- tests/unit/rabitqlib/index/qg_test.cpp | 5 +- tests/unit/rabitqlib/utils/space_test.cpp | 11 +- 9 files changed, 318 insertions(+), 8 deletions(-) create mode 100644 .github/scripts/repair_windows_wheel.py create mode 100644 .github/scripts/test_repair_windows_wheel.py diff --git a/.github/scripts/ci_changes.py b/.github/scripts/ci_changes.py index d05596a3..44837891 100644 --- a/.github/scripts/ci_changes.py +++ b/.github/scripts/ci_changes.py @@ -164,6 +164,11 @@ def classify(paths): selected.add("python_quality") elif path.startswith("sample/python/"): selected.update(("python_quality", "python_build", "wheels")) + elif path in ( + ".github/scripts/repair_windows_wheel.py", + ".github/scripts/test_repair_windows_wheel.py", + ): + selected.update(("python_quality", "wheels")) elif path.startswith(".github/scripts/"): selected.add("python_quality") elif path == ".github/workflows/test.yaml": diff --git a/.github/scripts/repair_windows_wheel.py b/.github/scripts/repair_windows_wheel.py new file mode 100644 index 00000000..547e1510 --- /dev/null +++ b/.github/scripts/repair_windows_wheel.py @@ -0,0 +1,152 @@ +"""Repair x64 wheels using VS runtimes, then reject older bundled runtimes.""" + +import argparse +import os +import subprocess +import sys +import zipfile +from pathlib import Path + +RUNTIME_PREFIXES = ("msvcp140", "vcruntime140", "vcomp140", "concrt140") + + +def pe_details(data): + # pefile is installed by the existing delvewheel before-build hook. + import pefile + + with pefile.PE(data=data) as image: + if image.FILE_HEADER.Machine != 0x8664: + raise RuntimeError("Windows wheel repair requires x64 PE binaries") + toolset = ( + image.OPTIONAL_HEADER.MajorLinkerVersion, + image.OPTIONAL_HEADER.MinorLinkerVersion, + ) + version = None + if getattr(image, "VS_FIXEDFILEINFO", None): + info = image.VS_FIXEDFILEINFO[0] + version = ( + info.FileVersionMS >> 16, + info.FileVersionMS & 0xFFFF, + info.FileVersionLS >> 16, + info.FileVersionLS & 0xFFFF, + ) + return toolset, version + + +def runtime_version(data): + _, version = pe_details(data) + if version is None: + raise RuntimeError("Runtime DLL has no file-version resource") + return version + + +def wheel_toolset(wheel): + with zipfile.ZipFile(wheel) as archive: + extensions = [ + name + for name in archive.namelist() + if Path(name).name.startswith("_rabitqlib.") and name.endswith(".pyd") + ] + if len(extensions) != 1: + raise RuntimeError("Expected exactly one _rabitqlib extension in the wheel") + toolset, _ = pe_details(archive.read(extensions[0])) + return toolset + + +def find_runtime_dirs(installation, minimum): + candidates = [] + root = installation / "VC" / "Redist" / "MSVC" + for crt in root.glob("*/x64/Microsoft.VC*.CRT"): + main_dll = crt / "msvcp140.dll" + if not main_dll.is_file(): + continue + version = runtime_version(main_dll.read_bytes()) + if version[:2] < minimum: + continue + for openmp in crt.parent.glob("Microsoft.VC*.OpenMP"): + omp_dll = openmp / "vcomp140.dll" + if ( + omp_dll.is_file() + and runtime_version(omp_dll.read_bytes())[:2] >= minimum + ): + candidates.append((version, crt, openmp)) + if not candidates: + raise RuntimeError(f"No x64 VS runtime directories match toolset {minimum}") + _, crt, openmp = max(candidates, key=lambda candidate: candidate[0]) + return crt, openmp + + +def verify_runtime_versions(wheel, minimum): + found = set() + with zipfile.ZipFile(wheel) as archive: + for name in archive.namelist(): + filename = Path(name).name.lower() + if not filename.endswith(".dll") or not filename.startswith( + RUNTIME_PREFIXES + ): + continue + version = runtime_version(archive.read(name)) + print(f"Bundled {name}: {'.'.join(map(str, version))}", flush=True) + if version[:2] < minimum: + raise RuntimeError( + f"Bundled {name} runtime {version} is older than toolset {minimum}" + ) + found.add(filename.split("-", 1)[0].removesuffix(".dll")) + if not {"msvcp140", "vcomp140"} <= found: + raise RuntimeError("Repaired wheel is missing its MSVC or OpenMP runtime") + + +def repair(wheel, destination, installation): + minimum = wheel_toolset(wheel) + runtime_dirs = find_runtime_dirs(installation, minimum) + print( + f"Wheel toolset: {minimum}; runtime search directories: {runtime_dirs}", + flush=True, + ) + subprocess.run( + [ + sys.executable, + "-m", + "delvewheel", + "repair", + "--add-path", + os.pathsep.join(map(str, runtime_dirs)), + "-w", + str(destination), + str(wheel), + ], + check=True, + ) + verify_runtime_versions(destination / wheel.name, minimum) + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("wheel", type=Path) + parser.add_argument("destination", type=Path) + args = parser.parse_args() + vswhere = ( + Path(os.environ["ProgramFiles(x86)"]) + / "Microsoft Visual Studio/Installer/vswhere.exe" + ) + installation = subprocess.check_output( + [ + str(vswhere), + "-latest", + "-products", + "*", + "-requires", + "Microsoft.VisualStudio.Component.VC.Tools.x86.x64", + "-property", + "installationPath", + "-utf8", + ], + encoding="utf-8-sig", + ).strip() + if not installation: + raise RuntimeError("No Visual Studio C++ installation found") + repair(args.wheel, args.destination, Path(installation)) + + +if __name__ == "__main__": + main() diff --git a/.github/scripts/test_ci_changes.py b/.github/scripts/test_ci_changes.py index f1e4213c..27764570 100644 --- a/.github/scripts/test_ci_changes.py +++ b/.github/scripts/test_ci_changes.py @@ -75,6 +75,11 @@ def test_known_support_paths_do_not_select_unrelated_builds(self): "python/utils/io.py": {"python_quality", "python_build", "wheels"}, "tests/core_consumer/run.cmake": {"cpp", "python_quality"}, "tests/fixtures/core_api_v1/consumer.cpp": {"cpp"}, + ".github/scripts/repair_windows_wheel.py": {"python_quality", "wheels"}, + ".github/scripts/test_repair_windows_wheel.py": { + "python_quality", + "wheels", + }, } for path, expected in cases.items(): with self.subTest(path=path): diff --git a/.github/scripts/test_ci_workflow_contracts.py b/.github/scripts/test_ci_workflow_contracts.py index 86813822..a2d1f417 100644 --- a/.github/scripts/test_ci_workflow_contracts.py +++ b/.github/scripts/test_ci_workflow_contracts.py @@ -24,7 +24,10 @@ def test_wheels_preserve_abi_coverage_and_release_gates(self): # Keep the full installed-wheel suite and platform repair hooks. self.assertIn("pytest {project}/tests/python", wheels["test-command"]) self.assertIn("delocate-wheel", wheels["macos"]["repair-wheel-command"]) - self.assertIn("delvewheel repair", wheels["windows"]["repair-wheel-command"]) + self.assertIn( + ".github/scripts/repair_windows_wheel.py", + wheels["windows"]["repair-wheel-command"], + ) self.assertIn("scripts/build-macos-openmp.sh", wheels["macos"]["before-all"]) self.assertEqual(workflow.count("pattern: '*-wheels'"), 2) artifacts = set() diff --git a/.github/scripts/test_repair_windows_wheel.py b/.github/scripts/test_repair_windows_wheel.py new file mode 100644 index 00000000..89d4f62f --- /dev/null +++ b/.github/scripts/test_repair_windows_wheel.py @@ -0,0 +1,137 @@ +"""Exercise runtime selection and fail-closed wheel repair without Windows tools.""" + +import json +import os +import subprocess +import tempfile +import unittest +import zipfile +from pathlib import Path +from unittest.mock import patch + +import repair_windows_wheel as repair + + +def metadata(toolset=(14, 51), version=None): + return json.dumps([toolset, version]).encode() + + +def read_metadata(data): + toolset, version = json.loads(data) + return tuple(toolset), tuple(version) if version is not None else None + + +class WindowsWheelRepairTest(unittest.TestCase): + def setUp(self): + temporary = tempfile.TemporaryDirectory() + self.addCleanup(temporary.cleanup) + self.root = Path(temporary.name) + self.installation = self.root / "Visual Studio" + mock = patch.object(repair, "pe_details", side_effect=read_metadata) + mock.start() + self.addCleanup(mock.stop) + + def runtimes(self, folder, version, arch="x64", openmp=True): + root = self.installation / "VC/Redist/MSVC" / folder / arch + crt = root / "Microsoft.VC143.CRT" + crt.mkdir(parents=True) + (crt / "msvcp140.dll").write_bytes(metadata(version=version)) + omp = root / "Microsoft.VC143.OpenMP" + if openmp: + omp.mkdir() + (omp / "vcomp140.dll").write_bytes(metadata(version=version)) + return crt, omp + + def wheel(self, path, dlls): + with zipfile.ZipFile(path, "w") as archive: + archive.writestr("rabitqlib/_rabitqlib.cp311-win_amd64.pyd", metadata()) + for name, version in dlls.items(): + archive.writestr("rabitqlib.libs/" + name, metadata(version=version)) + + def test_selects_matching_x64_runtime_not_older_or_wrong_architecture(self): + self.runtimes("14.40", (14, 40, 33810, 0)) + expected = self.runtimes("14.51", (14, 51, 36231, 0)) + self.runtimes("14.99", (14, 99, 0, 0), arch="x86") + self.assertEqual( + repair.find_runtime_dirs(self.installation, (14, 51)), expected + ) + + def test_checks_actual_file_version_instead_of_directory_name(self): + self.runtimes("14.99", (14, 40, 33810, 0)) + with self.assertRaisesRegex(RuntimeError, "No x64 VS runtime"): + repair.find_runtime_dirs(self.installation, (14, 51)) + + def test_requires_openmp_from_the_same_redist_tree(self): + self.runtimes("14.51", (14, 51, 36231, 0), openmp=False) + self.runtimes("14.40", (14, 40, 33810, 0)) + with self.assertRaisesRegex(RuntimeError, "No x64 VS runtime"): + repair.find_runtime_dirs(self.installation, (14, 51)) + + def test_rejects_observed_mixed_runtime_wheel(self): + wheel = self.root / "mixed.whl" + self.wheel( + wheel, + { + "msvcp140-old.dll": (14, 40, 33810, 0), + "msvcp140_1-new.dll": (14, 51, 36231, 0), + "vcomp140-new.dll": (14, 51, 36231, 0), + }, + ) + with self.assertRaisesRegex(RuntimeError, "older than toolset"): + repair.verify_runtime_versions(wheel, repair.wheel_toolset(wheel)) + + def test_rejects_missing_or_unversioned_runtime(self): + wheel = self.root / "missing.whl" + for dlls in ({}, {"msvcp140-hash.dll": (14, 51, 36231, 0)}): + self.wheel(wheel, dlls) + with self.assertRaisesRegex(RuntimeError, "missing"): + repair.verify_runtime_versions(wheel, (14, 51)) + self.wheel(wheel, {"msvcp140-hash.dll": None}) + with self.assertRaisesRegex(RuntimeError, "no file-version"): + repair.verify_runtime_versions(wheel, (14, 51)) + + def test_repair_prioritizes_selected_directories_and_checks_output(self): + paths = self.runtimes("14.51", (14, 51, 36231, 0)) + source = self.root / "source.whl" + destination = self.root / "output" + destination.mkdir() + self.wheel(source, {}) + + def run(command, *, check): + self.assertTrue(check) + self.assertEqual(command[1:4], ["-m", "delvewheel", "repair"]) + self.assertEqual( + command[command.index("--add-path") + 1], + os.pathsep.join(map(str, paths)), + ) + self.wheel( + destination / source.name, + { + "msvcp140-hash.dll": (14, 51, 36231, 0), + "vcomp140-hash.dll": (14, 51, 36231, 0), + }, + ) + + with patch.object(repair.subprocess, "run", side_effect=run): + repair.repair(source, destination, self.installation) + + def bad_output(*args, **kwargs): + self.wheel( + destination / source.name, {"msvcp140-old.dll": (14, 40, 33810, 0)} + ) + + with patch.object(repair.subprocess, "run", side_effect=bad_output): + with self.assertRaisesRegex(RuntimeError, "older than toolset"): + repair.repair(source, destination, self.installation) + + with patch.object( + repair.subprocess, + "run", + side_effect=subprocess.CalledProcessError(1, "repair"), + ): + with self.assertRaises(subprocess.CalledProcessError): + repair.repair(source, destination, self.installation) + + +if __name__ == "__main__": + unittest.main() diff --git a/pyproject.toml b/pyproject.toml index 408d768a..545326ba 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -75,7 +75,7 @@ CMAKE_ARGS = '-DRABITQ_ENABLE_NATIVE_OPTIMIZATION=OFF -DOpenMP_ROOT="$RABITQ_OPE [tool.cibuildwheel.windows] before-build = "python -m pip install delvewheel" -repair-wheel-command = "delvewheel repair -w {dest_dir} {wheel}" +repair-wheel-command = 'python "{project}/.github/scripts/repair_windows_wheel.py" "{wheel}" "{dest_dir}"' [tool.pytest.ini_options] testpaths = ["tests/python"] diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt index a8fa0441..c9d8d636 100644 --- a/tests/CMakeLists.txt +++ b/tests/CMakeLists.txt @@ -157,10 +157,12 @@ endif() # Run allocation replacement in a separate process; this also exercises core-only builds. add_executable(compact_codes_failure_test failures/compact_codes_failure_test.cpp) -target_link_libraries(compact_codes_failure_test PRIVATE GTest::gtest_main rabitqlib::core) +target_link_libraries(compact_codes_failure_test PRIVATE GTest::gtest_main) if(RABITQ_BUILD_FASTSCAN) target_link_libraries(compact_codes_failure_test PRIVATE rabitqlib::fastscan) target_compile_definitions(compact_codes_failure_test PRIVATE RABITQ_TEST_FASTSCAN) +else() + target_link_libraries(compact_codes_failure_test PRIVATE rabitqlib::core) endif() rabitq_configure_target(compact_codes_failure_test) if(CMAKE_SYSTEM_NAME STREQUAL "Linux") diff --git a/tests/unit/rabitqlib/index/qg_test.cpp b/tests/unit/rabitqlib/index/qg_test.cpp index ec254f4c..e20da63e 100644 --- a/tests/unit/rabitqlib/index/qg_test.cpp +++ b/tests/unit/rabitqlib/index/qg_test.cpp @@ -949,7 +949,10 @@ TEST(QGEstimatorTest, CandidateMaskBackendsMatchScalarOrdering) { distances[29] = -std::numeric_limits::infinity(); constexpr float kThreshold = 2.0F; - const uint32_t expected = simd::qg_candidate_mask_generic(distances.data(), kThreshold); + // Keep residues 0..4 in each group of seven, except +inf at lane 17. + // Unordered NaN at lane 3 and -inf at lane 29 remain candidates. + constexpr uint32_t expected = 0xF3E5CF9FU; + EXPECT_EQ(simd::qg_candidate_mask_generic(distances.data(), kThreshold), expected); #if defined(__x86_64__) || defined(_M_X64) if (cpu::has_avx2()) { EXPECT_EQ(simd::qg_candidate_mask_avx2(distances.data(), kThreshold), expected); diff --git a/tests/unit/rabitqlib/utils/space_test.cpp b/tests/unit/rabitqlib/utils/space_test.cpp index b423a401..19a6ce49 100644 --- a/tests/unit/rabitqlib/utils/space_test.cpp +++ b/tests/unit/rabitqlib/utils/space_test.cpp @@ -219,6 +219,9 @@ TEST(WarmupIpX0Q, SupportsUnalignedCodes) { const float expected = delta * static_cast(weighted_intersection) + vl * static_cast(data_popcount); + EXPECT_FLOAT_EQ( + warmup_ip_x0_q_512(codes, query.data(), delta, vl, dim, b_query), expected + ); #if defined(__x86_64__) || defined(_M_X64) if (cpu::has_avx2()) { EXPECT_FLOAT_EQ( @@ -227,10 +230,6 @@ TEST(WarmupIpX0Q, SupportsUnalignedCodes) { ), expected ); - EXPECT_FLOAT_EQ( - warmup_ip_x0_q_512(codes, query.data(), delta, vl, dim, b_query), - expected - ); } #endif #if defined(__x86_64__) || defined(_M_X64) @@ -250,6 +249,10 @@ TEST(WarmupIpX0Q, SupportsUnalignedCodes) { TEST(WarmupIpX0Q, RejectsQueryWidthsBeyondByte) { const std::array data{}; const std::array query{}; + EXPECT_THROW( + warmup_ip_x0_q_512(data.data(), query.data(), 1.0F, 0.0F, 64, 9), + std::invalid_argument + ); #if defined(__x86_64__) || defined(_M_X64) if (cpu::has_avx2()) { EXPECT_THROW(