Skip to content

Commit

Permalink
Enable Windows python bindings (#343)
Browse files Browse the repository at this point in the history
* Use int64 for counter to fix windows compilation error

* Fix windows python bindings by adding install_lib command to move windows build output into python package

* Update to use Path instead of os

* Change batch_insert num_inserts signature to signed type for OpenMP compatibility

* Update num_inserts to int32_t per PR request

---------

Co-authored-by: Nick Caurvina <[email protected]>
  • Loading branch information
nicaurvi and Nick Caurvina authored May 9, 2023
1 parent 3db6139 commit b011dcb
Show file tree
Hide file tree
Showing 3 changed files with 43 additions and 6 deletions.
2 changes: 2 additions & 0 deletions .gitignore
Original file line number Diff line number Diff line change
Expand Up @@ -371,3 +371,5 @@ wheelhouse/*
dist/*
venv*/**
*.swp

gperftools
4 changes: 2 additions & 2 deletions python/src/diskann_bindings.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -194,7 +194,7 @@ template <class T> struct DynamicInMemIndex
}

auto batch_insert(py::array_t<T, py::array::c_style | py::array::forcecast> &vectors,
py::array_t<IdT, py::array::c_style | py::array::forcecast> &ids, const size_t num_inserts,
py::array_t<IdT, py::array::c_style | py::array::forcecast> &ids, const int32_t num_inserts,
const int num_threads = 0)
{
if (num_threads == 0)
Expand All @@ -204,7 +204,7 @@ template <class T> struct DynamicInMemIndex
py::array_t<int> insert_retvals(num_inserts);

#pragma omp parallel for schedule(dynamic, 1)
for (size_t i = 0; i < num_inserts; i++)
for (int32_t i = 0; i < num_inserts; i++)
{
insert_retvals.mutable_data()[i] = _index->insert_point(vectors.data(i), *(ids.data(i)));
}
Expand Down
43 changes: 39 additions & 4 deletions setup.py
Original file line number Diff line number Diff line change
Expand Up @@ -3,12 +3,14 @@

import os
import re
import shutil
import subprocess
import sys
from pathlib import Path

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

# Convert distutils Windows platform specifiers to CMake -A arguments
PLAT_TO_CMAKE = {
Expand Down Expand Up @@ -44,7 +46,7 @@ def build_extension(self, ext: CMakeExtension) -> None:
f"-DCMAKE_LIBRARY_OUTPUT_DIRECTORY={extdir}{os.sep}",
f"-DPYTHON_EXECUTABLE={sys.executable}",
f"-DCMAKE_BUILD_TYPE={cfg}", # not used on MSVC, but no harm
f"-DVERSION_INFO={self.distribution.get_version()}" # commented out because we want this set in the CMake file
f"-DVERSION_INFO={self.distribution.get_version()}" # commented out, we want this set in the CMake file
]
build_args = []
# Adding CMake arguments set as environment variable
Expand Down Expand Up @@ -125,10 +127,43 @@ def build_extension(self, ext: CMakeExtension) -> None:
)


class InstallCMakeLibs(install_lib):
def run(self):
"""
Windows only copy from the x64/Release directory and place them in the package
"""

self.announce("Moving library files", level=3)

self.skip_build = True

# we only need to move the windows build output
windows_build_output_dir = Path('.') / 'x64' / 'Release'

if windows_build_output_dir.exists():
libs = [
os.path.join(windows_build_output_dir, _lib) for _lib in
os.listdir(windows_build_output_dir) if
os.path.isfile(os.path.join(windows_build_output_dir, _lib)) and
os.path.splitext(_lib)[1] in [".dll", '.lib', '.pyd', '.exp']
]

for lib in libs:
shutil.move(
lib,
os.path.join(self.build_dir, 'diskannpy', os.path.basename(lib))
)

super().run()


setup(
ext_modules=[CMakeExtension("diskannpy._diskannpy", ".")],
cmdclass={"build_ext": CMakeBuild},
cmdclass={
"build_ext": CMakeBuild,
'install_lib': InstallCMakeLibs
},
zip_safe=False,
package_dir = {"diskannpy": "python/src"},
package_dir={"diskannpy": "python/src"},
exclude_package_data={"diskannpy": ["diskann_bindings.cpp"]}
)

0 comments on commit b011dcb

Please sign in to comment.