From 01d636247218d8ae361d2d825c2da1bbc041ef7e Mon Sep 17 00:00:00 2001
From: Philipp Moritz <pcmoritz@gmail.com>
Date: Thu, 26 Sep 2019 12:55:14 -0700
Subject: [PATCH] Serialize StringIO with pickle (#5781)

---
 python/ray/tests/test_basic.py |  8 ++++++++
 python/ray/worker.py           | 11 +++++++++++
 2 files changed, 19 insertions(+)

diff --git a/python/ray/tests/test_basic.py b/python/ray/tests/test_basic.py
index 0381e3fa1..154f114bd 100644
--- a/python/ray/tests/test_basic.py
+++ b/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
diff --git a/python/ray/worker.py b/python/ray/worker.py
index b8d6f6e85..4a6c36ae0 100644
--- a/python/ray/worker.py
+++ b/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,