dask / dask/dask-ml

Dask Incremental with a slow fit will grow progressively slower due to spilling to disk

Open
#765 6 comments 1 reaction 0 assignees View on GitHub
Dominant language
Python
Stars
951
Forks
262
PR merge metrics
No merged PRs in 30d

Description

When using Dask incremental to train on datasets that are larger than memory if fit is substantially slower than the read tasks then most of the array will end up spilled back to disk. When fit hits those spilled tasks and has to load them this makes fit even slower and allows the read tasks that are reading data from disk to get further ahead making each fit operation spend time loading spilled data from disk.

Here is some example code that demonstrates this issue:
```python
from distributed.client import Client
import tensorflow as tf
from tensorflow.keras.layers import Dense
from tensorflow.keras.models import Sequential
from scikeras.wrappers import KerasClassifier
from tensorflow.keras.datasets import mnist
from tensorflow.keras.utils import to_categorical
import numpy as np
from typing import Tuple
import dask_ml
from dask import array
import dask
import sys
from dask_ml import datasets

if __name__ == "__main__":
X,y = datasets.make_classification_df(n_samples = int(sys.argv[1]),n_features=784,n_classes=2,chunks=(3000,784))

X['ys'] = y
X.to_parquet(sys.argv[2])
```
I used this script to generate data 4x the amount of RAM that I have.

Then I trained on it using incremental:
```python
import tensorflow as tf
from tensorflow.keras.layers import Dense
from tensorflow.keras.models import Sequential
from scikeras.wrappers import KerasClassifier
from tensorflow.keras.datasets import mnist
from tensorflow.keras.utils import to_categorical
import numpy as np
from typing import Tuple
import dask_ml
from dask import array
from dask import dataframe
from dask.distributed import LocalCluster
import dask
import sys

def build_model(lr=0.01, momentum=0.9):

layers = [Dense(512, input_shape=(784,), activation="relu"),
Dense(10, input_shape=(512,), activation="softmax")]
model = Sequential(layers)

opt = tf.keras.optimizers.SGD(
learning_rate=lr, momentum=momentum, nesterov=True,
)
model.compile(loss="binary_crossentropy", optimizer=opt, metrics=['accuracy'])
return model

if __name__ == "__main__":
cluster = LocalCluster(n_workers=4,memory_target_fraction=.4, memory_spill_fraction=.4)
dask.distributed.Client(cluster)

path = sys.argv[1]

model = KerasClassifier(build_fn=build_model, lr=0.1, momentum=0.9,batch_size=1)
inc_mod = dask_ml.wrappers.Incremental(model)

df = dataframe.read_parquet(path)

parition_sizes = list(df.map_partitions(lambda x: x.shape[0]).compute())
print(type(parition_sizes))

Xs = df.loc[:,df.columns != 'ys'].to_dask_array(lengths=parition_sizes)

ys = df.loc[:,'ys'].to_dask_array(lengths=parition_sizes)

inc_mod.fit(Xs,ys)
```
I dropped the batch size to 1 to intentionally make fit slower and dropped the spill fractions to make the issue occur faster.

This results in this computation graph:
![big_graph](https://user-images.githubusercontent.com/19430893/101300287-19922e80-3803-11eb-92ca-38c03a6c3e54.png)

A more zoomed in version:
![small_graph](https://user-images.githubusercontent.com/19430893/101300297-2151d300-3803-11eb-9a98-365e060627fd.png)

As you can see the 3 early tasks which are reads and transformations are all dependencies of each fit task while the fit tasks are each dependent on the previous fit task in order to guarantee training the model in serial with respect to dask.

The trouble is dask will continue to execute the dependencies of each fit ad infinium
![2020-12-06-201720_954x281_scrot](https://user-images.githubusercontent.com/19430893/101300382-6970f580-3803-11eb-98f0-d72707dccbe4.png)

When fit is slow these reads will get far ahead of the fit task. Since incremental is meant for larger than memory datasets eventually the tasks will start spilling to disk. Fit will then have to load the dependency back from disk when it reaches the spilled block. These loads can be quite slow if blocks are large. In the worst case this ends up as essentially serial execution where each fit is preceded by a block read. All of these reads and writes can make training extremely slow as spilling, fit loading spilled tasks, and further read tasks are all competing for disk time.

Ideally what would occur is read tasks are only executed until memory is full and then workers wait until fit is done with some tasks freeing up space to read in more training data. That way data never has to be spilled to disk and read back.

One way I thought about acheiving this is by modifying the graph to add synchronization points at regular intervals.

For example if we have the task graph [read0->fit0, read1->fit1,read2->fit2,read3->fit3,read4->fit4, read5->fit5,read6->fit6,read7->fit7, read8->fit8, fit0->fit1,fit1->fit2, fit2->fit3,fit3->fit4, fit4->fit5,fit5->fit6, fit6->fit7,fit7->fit8]

And we know we can fit 3 blocks in memory then we can add edges {fit2->read3, fit5->read6}. These artificial dependencies force the reading tasks to wait until fit has released the 3 blocks before starting to read the next set of blocks. This behavior is not ideal because it will wait to start reading the next 3 chunks until all 3 previous chunks have been read. Ideally as soon as the first chunk is fit and released we can start loading the fourth chunk. That being said I do not believe such behavior can easily be accomplished with normal computation graphs and would likely require using raw futures.

I think this is an important issue for Incremental since its purpose is larger than memory datasets and it is not super uncommon for model fitting to take longer than reading and transforming the data.

Contributor guide

Open the contributing guide

Research direction

Start at dask_ml.wrappers.Incremental and reproduce the provided LocalCluster, parquet, and KerasClassifier example with a dataset larger than memory. Inspect the resulting read and fit task dependencies and worker spilling behavior. Done means slow incremental fitting no longer allows reads to run far ahead and repeatedly spill training blocks to disk.

Written by the indexing model from the issue text.

Assessment

Tech stack
python, tensorflow
Domain
distributed-systems, machine-learning, performance
Issue type
Bug
Difficulty
5/5
Estimated time
Over a week
Activity status
Stale
Clarity
Mostly clear
Newbie friendliness
30/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.