kammerje / spaceKLIP

Pipeline for reducing JWST high-contrast imaging data. Published in Kammerer et al. 2022 and Carter et al. 2022.
https://ui.adsabs.harvard.edu/abs/2022SPIE12180E..3NK/abstract
MIT License
16 stars 9 forks source link

Apple Silcone (M1/M2/M3) chips and pyklip.parallelized multiprocessing #169

Open wbalmer opened 2 months ago

wbalmer commented 2 months ago

As discussed in the pyklip documentation, there are issues with the multiprocessing interactions for numpy compiled with BLAS or MKL. In my own experience these only cropped up when I upgraded computers to an Apple Silicone M3 chip, because conda for non-intel chips compiles numpy with OPENBLAS and not MKL. Apparently there are also MKL problems similar to this for some compilations of numpy.

In my case, my M3 chip significantly sped up all my postprocessing steps wrapped by spaceklip, including the webbpsf computations (great!) but would stall out and steal all my CPU when running klip itself. As per the pyklip docs and Jason's suggestion, I need to add the following lines to my code block before calling pyklippipeline.run_obs but after doing my preprocessing, e.g.

os.environ["OPENBLAS_NUM_THREADS"] = "1"

the only way to detect my numpy compilation was OPENBLAS and not MKL is np.show_config() but this returns NoneType and can't be checked. My solution (pasted below) was to use platform to check if the system processor is "arm" or intel but as explained by Jason this is not a good solution since you could conceivably have an arm processor with some frankenstein MKL numpy or an intel processor with an OPENBLAS numpy.

import platform
if platform.processor() == 'arm':
    os.environ["OPENBLAS_NUM_THREADS"] = "1"
pyklippipeline.run_obs(database=Database,

Maybe worth some warning in spaceklip tutorial/documentation or printed when pyklippipeline.run_obs is called?