Serialize StringIO with pickle (#5781)

python/ray/tests/test_basic.py

@@ -6,6 +6,7 @@ from __future__ import print_function
 import collections
 from concurrent.futures import ThreadPoolExecutor
 import glob
+import io
 import json
 import logging
 from multiprocessing import Process
@@ -304,6 +305,13 @@ def test_complex_serialization(ray_start_regular):
         assert_equal(obj, ray.get(f.remote(obj)))
         assert_equal(obj, ray.get(ray.put(obj)))
 
+    # Test StringIO serialization
+    s = io.StringIO(u"Hello, world!\n")
+    s.seek(0)
+    line = s.readline()
+    s.seek(0)
+    assert ray.get(ray.put(s)).readline() == line
+
 
 def test_nested_functions(ray_start_regular):
     # Make sure that remote functions can use other values that are defined

python/ray/worker.py

@@ -8,6 +8,7 @@ import atexit
 import faulthandler
 import hashlib
 import inspect
+import io
 import json
 import logging
 import numpy as np
@@ -1278,6 +1279,16 @@ def _initialize_serialization(job_id, worker=global_worker):
             local=True,
             job_id=job_id,
             class_id="ray.signature.FunctionSignature")
+        # Tell Ray to serialize StringIO with pickle. We do this because
+        # Ray's default __dict__ serialization is incorrect for this type
+        # (the object's __dict__ is empty and therefore doesn't
+        # contain the full state of the object).
+        register_custom_serializer(
+            io.StringIO,
+            use_pickle=True,
+            local=True,
+            job_id=job_id,
+            class_id="io.StringIO")
 
 
 def init(address=None,