search

lmcinnes / umap

Uniform Manifold Approximation and Projection
BSD 3-Clause "New" or "Revised" License
7.19k stars 780 forks source link

numba AOT distribution #1051

Open pavlin-policar opened 10 months ago

pavlin-policar commented 10 months ago

umap-learn uses numba for JIT. This is fine for making the code run fast, however, this also makes distributing the umap-learn package a pain in other open-source software. We are interested in integrating UMAP into Orange, but we are hesitant to do so because of the numba and llvmlite dependencies, as this would add another 30Mb size to our installers, which are already too big.

I have recently learned that numba also has an Ahead-of-Time (AOT) mode, which precompiles binaries without any changes to the code. I haven't looked into it, but my understanding is that these binaries then do not require neither numba nor llvmlite to be installed. This would likely complicate the build and distributing process for umap-learn, since, if you went this route, you'd likely need to setup a CI system to build on all the major platforms and Python versions.

I am wondering if you have any plans to migrate to numba AOT, or distributing precompiled binaries in any other capacity?

lmcinnes commented 10 months ago

I think there is some potential benefit to providing binary wheels. A catch is that compilation specific to the processors of the host machine can be pretty beneficial, and binary wheels tend to be generic. That gives me some hesitancy of trying to distribute AOT based binary wheels directly on PyPI. On the other hand potentially a specific AOT based version of these libraries might be practical under a different name on PyPI. Alternatively if you are packing things up yourself you can build AOT versions for distribution? I'm open to ideas on this as I agree it would potentially be useful.

pavlin-policar commented 10 months ago

Yes, the binary wheels are more generic and typically won't be as fast than if they are compiled on the native system. I don't know how much numba is affected by this, but it isn't that drastic when compiling cython extensions. Noticeable, yes, but not drastic.

Even if you do start packaging binaries, you could still make a source distribution that would either install the non-AOT-ed version or it would perform the AOT compilation on the host machine. Users can opt for the sdist version with the --no-binary flag in pip, which installs packages from source. However, users will have to remember to use it because wheels get installed by default. Releasing a new, separate package would solve these issues, and seems like a good solution.

I haven't played around with this at all, so I can't really comment on setting up numba's AOT. But looking at the documentation, it seems pretty straightforward. We could probably build the umap binaries for orange ourselves, but we tend to try to avoid that since it just complicates and already complicated build pipeline.

lmcinnes commented 10 months ago

I'll try to look into the AOT a little more and see if I can have umap-learn-aot and pynndescent-aot packages available on PyPI; hopefully that would resolve some of the issues for you.

pavlin-policar commented 10 months ago

Thanks, that would be very helpful!

ritvikshrivastava commented 3 months ago

Hi @lmcinnes, found this thread to be extremely resonating in terms of the concerns highlighted by @pavlin-policar . I was wondering if you reached any resolution for the AOT packages and publishing the binary wheels? Thanks!

lmcinnes commented 3 months ago

No further news I am afraid. I'm really dependent on numba's AOT compilation to make this possible. I understand there is work being done on this, but it is not yet really in a state where we can distribute binary packages.