lgrcia
Update periodic search
When running the periodic search I repeatedly got the warning: "os.fork() is incompatible with multithreaded code, and JAX is multithreaded, so this will likely lead to a deadlock."
It never seemed to cause any actual issues, but fork is set to be deprecated so I modified the periodic search to use 'spawn' instead. Since 'spawn' (and 'forkserver') don't allow the use of global fold_f I had to switch to Pool.starmap and explicitly pass fold_f.
While I was at it I also exposed processes and chunksize to allow the user more fine-grained control over the multiprocessing and added if __name__ == '__main__': which I understand to be good practice for multiprocessed code.
As a final note I wonder if it might be a good idea to move the snr_f evaluations outside of the with Pool statement ?
Thanks for this fix @talensgj. I am curious, did you see any performance benefits from setting the chunksize?
It was hard to tell. In the original use with imap with chunksize=1 both multiprocessing and jax were running multi-core processing at the same time. With starmap and chunksize>1 this no longer happens.
I would need to compare chunksize=1 and chunksize=500 on the new code to check, but I've been distracted by other things.
I can at least confirm that with chunksize=500, about 30 mins were spend on solve_ with multiprocessing and 90 mins on snr_f with jax.