[tune] Allow fetching pinned objects from trainable functions (#1895)

* updates

* lint

* Update util.py

* Update function_runner.py

* updates

python/ray/tune/util.py

@@ -3,15 +3,25 @@ from __future__ import division
 from __future__ import print_function
 
 import base64
+import queue
+import threading
 
 import ray
 from ray.tune.registry import _to_pinnable, _from_pinnable
 
 _pinned_objects = []
+_fetch_requests = queue.Queue()
 PINNED_OBJECT_PREFIX = "ray.tune.PinnedObject:"
 
 
 def pin_in_object_store(obj):
+    """Pin an object in the object store.
+
+    It will be available as long as the pinning process is alive. The pinned
+    object can be retrieved by calling get_pinned_object on the identifier
+    returned by this call.
+    """
+
     obj_id = ray.put(_to_pinnable(obj))
     _pinned_objects.append(ray.get(obj_id))
     return "{}{}".format(PINNED_OBJECT_PREFIX,
@@ -19,12 +29,40 @@ def pin_in_object_store(obj):
 
 
 def get_pinned_object(pinned_id):
+    """Retrieve a pinned object from the object store."""
+
     from ray.local_scheduler import ObjectID
+
+    if threading.current_thread().getName() != "MainThread":
+        placeholder = queue.Queue()
+        _fetch_requests.put((placeholder, pinned_id))
+        print("Requesting main thread to fetch pinned object", pinned_id)
+        return placeholder.get()
+
     return _from_pinnable(
         ray.get(
             ObjectID(base64.b64decode(pinned_id[len(PINNED_OBJECT_PREFIX):]))))
 
 
+def _serve_get_pin_requests():
+    """This is hack to avoid ray.get() on the function runner thread.
+
+    The issue is that we run trainable functions on a separate thread,
+    which cannot access Ray API methods. So instead, that thread puts the
+    fetch in a queue that is periodically checked from the main thread.
+    """
+
+    assert threading.current_thread().getName() == "MainThread"
+
+    try:
+        while not _fetch_requests.empty():
+            (placeholder, pinned_id) = _fetch_requests.get_nowait()
+            print("Fetching pinned object from main thread", pinned_id)
+            placeholder.put(get_pinned_object(pinned_id))
+    except queue.Empty:
+        pass
+
+
 if __name__ == '__main__':
     ray.init()
     X = pin_in_object_store("hello")