Ray Tensorflow-gpu 2.0 RecursionError

Основная проблема, по-видимому, заключается в том, что cloudpickle (который Ray использует для сериализации удаленных функций и отправки их рабочим процессам) не может обработать класс tf.keras.Sequential. Например, я могу воспроизвести проблему следующим образом

import cloudpickle  # cloudpickle.__version__ == '1.2.1'
import tensorflow as tf  # tf.__version__ == '2.0.0-rc0'

def f():
    tf.keras.Sequential

cloudpickle.loads(cloudpickle.dumps(f))  # This fails.

Последняя строка не работает с

---------------------------------------------------------------------------
RecursionError                            Traceback (most recent call last)
<ipython-input-23-25cc307e6227> in <module>
----> 1 cloudpickle.loads(cloudpickle.dumps(f))

~/anaconda3/lib/python3.6/site-packages/tensorflow/__init__.py in __getattr__(self, item)
     48 
     49   def __getattr__(self, item):
---> 50     module = self._load()
     51     return getattr(module, item)
     52 

~/anaconda3/lib/python3.6/site-packages/tensorflow/__init__.py in _load(self)
     42   def _load(self):
     43     """Import the target module and insert it into the parent's namespace."""
---> 44     module = _importlib.import_module(self.__name__)
     45     self._parent_module_globals[self._local_name] = module
     46     self.__dict__.update(module.__dict__)

... last 2 frames repeated, from the frame below ...

~/anaconda3/lib/python3.6/site-packages/tensorflow/__init__.py in __getattr__(self, item)
     48 
     49   def __getattr__(self, item):
---> 50     module = self._load()
     51     return getattr(module, item)
     52 

RecursionError: maximum recursion depth exceeded while calling a Python object

Интересно, что это успешно с tensorflow==1.14.0, но я полагаю, что keras сильно изменился в 2.0.

Обходной путь

В качестве обходного пути вы можете попробовать определить f в отдельном модуле или файле Python, например

# helper_file.py

import tensorflow as tf

def f():
    tf.keras.Sequential

А затем используйте его в своем основном сценарии следующим образом.

import helper_file
import ray

ray.init(num_cpus=1)

@ray.remote
def use_f():
    helper_file.f()

ray.get(use_f.remote())

Разница здесь в том, что когда cloudpickle пытается сериализовать use_f, он фактически не смотрит на содержимое helper_file. Когда какой-либо рабочий процесс пытается десериализовать use_f, этот рабочий процесс импортирует helper_file. Это дополнительное косвенное обращение, кажется, заставляет облачко работать более надежно. То же самое происходит, когда вы обрабатываете функцию, используя тензорный поток или любую библиотеку. Cloudpickle не сериализует всю библиотеку, он просто сообщает процессу десериализации импортировать соответствующую библиотеку.

Примечание. Чтобы это работало на нескольких машинах, helper_file.py должен существовать и находиться на пути Python на каждой машине (один из способов добиться этого - установить его как модуль Python на каждом компьютере).

Я подтвердил, что это, похоже, решает проблему в вашем примере. После исправления я столкнулся с

  File "<ipython-input-4-bb51dc74442c>", line 3, in play
  File "/Users/rkn/Workspace/ray/helper_file.py", line 15, in play
    action = actor.predict_classes(obs).flatten()[0]
AttributeError: 'NoneType' object has no attribute 'predict_classes'

но это похоже на отдельную тему.

См. Ответ GitHub на эту проблему: https://github.com/ray-project/ray/issues/5614

Все, что нужно сделать, это импортировать тензорный поток в определение функции:

@ray.remote
def play(weights):
    import tensorflow as tf
    actor = tf.keras.Sequential([
        tf.keras.layers.Dense(64, input_shape=(24,), activation='relu'),
        tf.keras.layers.Dense(4, activation='softmax')
        ])
    actor.set_weights(weights)
    env = gym.make('BipedalWalker-v2').env
    env._max_episode_steps=1e20
    obs = env.reset()
    for _ in range(1200):
        action = actor.predict_classes(np.array([obs])).flatten()[0]
        action = env.action_space.sample()
        obs, rt, done, info = env.step(action)
    return rt