k-means Memory usage about dask-ml

[TomAugspurger]

@mrocklin I don't see "failed to deserialize" error every time. Adding

diff --git a/distributed/protocol/serialize.py b/distributed/protocol/serialize.py
index d657dab..b7b492e 100644
--- a/distributed/protocol/serialize.py
+++ b/distributed/protocol/serialize.py
@@ -87,6 +87,8 @@ def typename(typ):
 def _find_lazy_registration(typename):
     toplevel, _, _ = typename.partition('.')
     if toplevel in lazy_registrations:
+        import time
+        time.sleep(.1)
         lazy_registrations.pop(toplevel)()
         return True
     else:

to serialize.py makes it consistent. Perhaps a race condition in the lazy importer? I'll see if I can narrow down the example further.

[TomAugspurger]

Collecting some some debugging observations:

The getitem takes ~0.01s normally with an indexer of length 1. There's no increase in memory usage. When things slow down, things take ~.5s. The slower tasks include a disk-read-getitem:

(slow on the left, normal on the right)

When slicing with multiple (e.g. size=2) the graphs sometimes look different. Fast:

slow:

Not sure if this is meaningful or not. I suspect it is. I assumed the order of operations would be

getitem for each worker -> concatenate results.

But if it's

transfer blocks to single worker -> getitem

that would explain the slowdown and memory increase.

[mrocklin]

disk-read- time blocks are due to getting elements out of worker.data. This could mean that there are many elements in worker.data that are in memory, or more likely that there are a few elements that are stored on disk. This also corresponds to colored bars in the upper left memory use plot in the scheduler.

[mrocklin]

If your getitem/transfer vs transfer/getitem question refers to the above code:

for i in range(5):
    idx = np.random.randint(0, len(X), size=5)
    centers = X[idx].compute()
    mem()

Then you should be fine. Dask.array definitely does the intelligent thing here.

[TomAugspurger]

Interestingly, whether or not the indexer is sorted seems to matter. Adding a sorted(idx) before indexing:

centers = c.compute(X[sorted(idx)])

outputs:

3.84 GB
3.84 GB
3.84 GB
3.84 GB
3.84 GB

Trying it out in k_init now.

[mrocklin]

There is a special fast-path within dask/array/slicing.py for when the index is sorted. You might search that file for issorted to find what the other non-fast-path is doing

[mrocklin]

Or just always sort

[TomAugspurger]

Just sorting works for me here. I'll still take a look in slicing to see if anything weird is going on.

[TomAugspurger]

Thanks!