Skip to content
New issue

Have a question about this project? # for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “#”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? # to your account

Update setup.py to build object files in parallel if requested #105

Open
wants to merge 2 commits into
base: master
Choose a base branch
from

Conversation

althonos
Copy link
Contributor

@althonos althonos commented May 27, 2024

Hi again!

Since the best way to install cython-blis is to compile it from source to take advantage of the machine architecture. In the case of our HPC cluster, I end up re-installing cython-blis on each node executor at the start of each job to make sure I'm using optimized code, but this takes a bit of time.

Given that BLIS has a lot of source files, the build process can be parallelized easily. I just changed the logic of the ExtensionBuilder.compile_objects code to actually invoke the compiler to build objects in parallel with a ThreadPool, based on the parallel flag of the command line (which is a default build_ext option), or using the MAX_JOBS environment variable (similar to what torch and flash-attn are doing).

By default, I left the job count to 1, so that parallel compilation happens only if enabled. Using 4 threads, the compilation is about twice faster:

MAX_JOBS="4" pip install blis --no-binary=blis 

# for free to join this conversation on GitHub. Already have an account? # to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

1 participant