dask_example.py

import numpy as np
import tensorflow as tf
import keras
from keras import backend as K
from distributed import Client, client


def new_session():
    sess = tf.Session()
    return sess


def test_simulation(size, sess):
    K.set_session(sess)

    X = np.random.random((10, 5))
    y = np.random.random((10, 1))
    model_input = keras.layers.Input((5,))
    hidden = keras.layers.Dense(size)(model_input)
    model_output = keras.layers.Dense(1)(hidden)

    model = keras.models.Model(model_input, model_output)
    model.compile(optimizer='sgd', loss='mse')
    model.fit(X, y, nb_epoch=10)


if __name__ == '__main__':
    SCHEDULER_IP = 'tcp://127.0.0.1:8786'
#    SCHEDULER_IP = '169.229.219.178:8786'
    client = Client(SCHEDULER_IP)
    
    sess = client.submit(new_session)
    res_list = [client.submit(test_simulation, i, sess) for i in range(5, 10)]
#    client.wait(res_list)

First attempt. Not effective, just recording for posterity:

import numpy as np
import tensorflow as tf
import keras
from keras import backend as K
from dask.utils import SerializableLock
from distributed import Client, LocalCluster, wait
from distributed.worker_client import get_worker

lock = SerializableLock('tensorflow-session')


def get_session():
    worker = get_worker()
    if not hasattr(worker, 'tensorflow_session'):
        worker.tensorflow_session = tf.Session()
    return worker.tensorflow_session


def test_simulation(size):
    with lock:  # this lock doesn't appear to help at all
        sess = get_session()
        K.set_session(sess)

    X = np.random.random((10, 5)) 
    y = np.random.random((10, 1)) 
    model_input = keras.layers.Input((5,))
    hidden = keras.layers.Dense(size)(model_input)
    model_output = keras.layers.Dense(1)(hidden)

    model = keras.models.Model(model_input, model_output)
    model.compile(optimizer='sgd', loss='mse')
    model.fit(X, y, nb_epoch=10)


if __name__ == '__main__':
    cluster = LocalCluster(n_workers=1, threads_per_worker=4, nanny=False)
    client = Client(cluster)

    res_list = [client.submit(test_simulation, i) for i in range(5, 10)]

    for res in res_list:
        try:
            print(res.result())
        except Exception as e:
            print(e)

Here is an example that doesn't use Dask at all that fails. I think that we should post this as a StackOverflow question:

from concurrent.futures import ThreadPoolExecutor

import numpy as np
import tensorflow as tf
import keras
from keras import backend as K

sess = tf.Session()


def test_simulation(size):
    K.set_session(sess)

    X = np.random.random((10, 5)) 
    y = np.random.random((10, 1)) 
    model_input = keras.layers.Input((5,))
    hidden = keras.layers.Dense(size)(model_input)
    model_output = keras.layers.Dense(1)(hidden)

    model = keras.models.Model(model_input, model_output)
    model.compile(optimizer='sgd', loss='mse')
    model.fit(X, y, nb_epoch=10)


if __name__ == '__main__':
    e = ThreadPoolExecutor(4)
    res_list = [e.submit(test_simulation, i) for i in range(5, 10)]

    for res in res_list:
        try:
            print(res.result())
        except Exception as e:
            print(e)

Output

Using TensorFlow backend.
/home/mrocklin/Software/anaconda/lib/python3.5/site-packages/keras/backend/tensorflow_backend.py:162: UserWarning: The default TensorFlow graph is not the graph associated with the TensorFlow session currently registered with Keras, and as such Keras was not able to automatically initialize a variable. You should consider registering the proper session with Keras via `K.set_session(sess)`.
  warnings.warn('The default TensorFlow graph is not the graph '
Epoch 1/10
Epoch 1/10
Cannot interpret feed_dict key as Tensor: Tensor Tensor("input_1:0", shape=(?, 5), dtype=float32) is not an element of this graph.
Cannot interpret feed_dict key as Tensor: Tensor Tensor("input_2:0", shape=(?, 5), dtype=float32) is not an element of this graph.
Epoch 1/10
Epoch 1/10
Cannot interpret feed_dict key as Tensor: Tensor Tensor("input_3:0", shape=(?, 5), dtype=float32) is not an element of this graph.
Cannot interpret feed_dict key as Tensor: Tensor Tensor("input_4:0", shape=(?, 5), dtype=float32) is not an element of this graph.
Epoch 1/10
Cannot interpret feed_dict key as Tensor: Tensor Tensor("input_5:0", shape=(?, 5), dtype=float32) is not an element of this graph.

The same seems to happens no matter where we set the session

You can even throw a lock around the entire contents of test_simulation and the system fails.

Any updated news on this? I am having the same issue.