I am training a neural network with a large text corpora. Each text generate quite a big matrix because I'm using a convolutional model. As my data won't feet in my still large memory, I try to stream it, and use keras.models fit_generator.
To feed keras, I have a pipeline composed of different preprocessing steps, that I arrange with a dask bag with lots of partitions. The dask bag reads a file on disk.
Even is dask is not handling iteration in a smart way (it just compute() and iter on result, which in my case blow up memory), I was to use something like this:
def compute_partition_iter(collection, **kwargs):
"""A utility to compute a collection items after items
"""
get = kwargs.pop("get", None) or _globals['get']
if get is None:
get = collection.__dask_scheduler__
postcompute_func, postcompute_args = collection.__dask_postcompute__()
dsk = collection.__dask_graph__()
for key in collection.__dask_keys__():
yield from f([partition], *args)
This compute partitions one by one and return items, computing next partition as we cross partition border.
This approach has a problem : it's only when we hit last item from partition that we provoque the computation of next elements, leading to a lag until next element. Within this lag, keras is stalled and we loose precious time !
So I imagine running the above compute_partition_iter
in a separate process thanks to multiprocessing.Pool
, feeding partitions in a Queue
with say 2 slots, so that in the generator, I won't always have one more partition ready.
But it seems that this is not supported by dask.bag
. I didn't dive deeply enough in the code, but it seems like there are some async methods used, or I don't know what.
Here is a reproductible code for the problem.
First a code that work, using a simple range.
import multiprocessing
import time
def put_q(n, q):
for i in range(n):
print(i, "<-")
q.put(i)
q.put(None)
q = multiprocessing.Queue(2)
with multiprocessing.Pool(1, put_q, (4, q)) as pool:
i = True
while i is not None:
print("zzz")
time.sleep(.5)
i = q.get()
if i is None:
break
print("-> ", i)
This outputs
0 <-
1 <-
2 <-
zzz
3 <-
-> 0
zzz
-> 1
zzz
-> 2
zzz
-> 3
zzz
you can see that, as expected, elements where computed in anticipation and it's all ok.
Now let's replace the range by a dask.bag
:
import multiprocessing
import time
import dask.bag
def put_q(n, q):
for i in dask.bag.from_sequence(range(n), npartitions=2):
print(i, "<-")
q.put(i)
q.put(None)
q = multiprocessing.Queue(5)
with multiprocessing.Pool(1, put_q, (4, q)) as pool:
i = True
while i is not None:
print("zzz")
time.sleep(.5)
i = q.get()
if i is None:
break
print("-> ", i)
In a jupyter notebook, it indefinitely raises :
Process ForkPoolWorker-71:
Traceback (most recent call last):
File "/usr/lib/python3.5/multiprocessing/process.py", line 249, in _bootstrap
self.run()
File "/usr/lib/python3.5/multiprocessing/process.py", line 93, in run
self._target(*self._args, **self._kwargs)
File "/usr/lib/python3.5/multiprocessing/pool.py", line 103, in worker
initializer(*initargs)
File "<ipython-input-3-e1e9ef9354a0>", line 8, in put_q
for i in dask.bag.from_sequence(range(n), npartitions=2):
File "/usr/local/lib/python3.5/dist-packages/dask/bag/core.py", line 1190, in __iter__
return iter(self.compute())
File "/usr/local/lib/python3.5/dist-packages/dask/base.py", line 154, in compute
(result,) = compute(self, traverse=False, **kwargs)
File "/usr/local/lib/python3.5/dist-packages/dask/base.py", line 407, in compute
results = get(dsk, keys, **kwargs)
File "/usr/local/lib/python3.5/dist-packages/dask/multiprocessing.py", line 152, in get
initializer=initialize_worker_process)
File "/usr/lib/python3.5/multiprocessing/context.py", line 118, in Pool
context=self.get_context())
File "/usr/lib/python3.5/multiprocessing/pool.py", line 168, in __init__
self._repopulate_pool()
File "/usr/lib/python3.5/multiprocessing/pool.py", line 233, in _repopulate_pool
w.start()
File "/usr/lib/python3.5/multiprocessing/process.py", line 103, in start
'daemonic processes are not allowed to have children'
AssertionError: daemonic processes are not allowed to have children
while the main process is stalled, waiting for elements in queue.
I also tried using a ipyparallel cluster but in this case the main process is simply stalled (no trace of the exception).
Does anyone knows the right way to do that ?
Is there a way I can run scheduler.get in parallel to my main code ?