I've been finding that the Dask arrays generated by this package don't give me the speed I would hope for, and this PR helps address that. (I'm working on another, too.)

My use case right now (my first with DKIST) is taking a ViSP data set and summing over wavelength to get a quasi-broadband view of what ViSP was seeing.

Doing this serially takes 1 m 53 s

ds = dkist.load_dataset("/home/svankooten/BXXDZ/")

def comp_slice(file):
    data = fits.getdata("/home/svankooten/BXXDZ/" + file)
    data = data[0]
    return data.sum(axis=0)

image = np.stack(list(map(comp_slice, ds.files.filenames)))

Doing this with multiprocessing takes 16 s with 8 cores

from multiprocessing import Pool

with Pool(8) as p:
    image = np.stack(p.map(comp_slice, ds.files.filenames))

The default dask scheduler uses threads, and with the GIL I get 250-300% CPU usage and it takes 1m 5s

ds.data.sum(axis=-2).compute();

Switching to Dask's distributed scheduler evades the GIL and ideally will look like the multiprocessing case, but it actually takes longer than the serial case, at 2m 33s

from dask.distributed import Client
client = Client(threads_per_worker=1, n_workers=8)
ds.data.sum(axis=-2).compute()

The task stream in the scheduler shows that there are lots of gaps between tasks

What I finally found as the cause is that each AstropyFITSLoader keeps a reference to its parent StripedExternalArray. As I understand it, the loaders get pickled and sent to the worker processes before they're used, and this reference causes these pickle payloads to be several megabytes each. This PR avoids pickling the StripedExternalArray. This makes the task stream much denser, and lets the computation complete in 47 s

(The x scale is different in this screenshot)

This beats the serial case (!) and the default threaded scheduler, but doesn't match multiprocessing. (Not sure how close this can get, I assume Dask has more overhead.) This PR doesn't change the run time for the default scheduler, which I suppose doesn't need to do any pickling.

This PR has the Loaders just record the basepath when they're created. If Filemanager.basepath is updated later on, the setter in StripedExternalArray now iterates through loader_array and updates each one. This has the downside of no longer having a single source of truth for basepath, but I think it's worth the speedup. (And it seems like a user would have to dig deep in internal API to change one but not the others.)

An alternate design I tried was to have have BaseFITSLoader keep its array_container reference, but I overrode BaseFITSLoader's pickling behavior so that only array_container.basepath is sent, not array_container itself. Then the unpickled loaders in the worker processes set a _basepath value (a snapshot of the "live" value at the time of pickling, which is when the dask calculation is triggered) but don't set array_container, and the basepath property returns whichever is set. This lets basepath be stored in only one location in the main process, but I think it has the side-effect that if a user loads a dataset, pickles it for some reason, unpickles it, updates the basepath and then triggers the Dask calculation, the pickling will have caused the loaders to drop their references and they won't get the updated basepath.

Hopefully this PR doesn't conflict with any future plans for FITSLoader!