Added thread-safe broadcast pickle registry

BryanCutler · BryanCutler · commit b703f8359cee · 2017-07-20T14:36:09.000-07:00
diff --git a/python/pyspark/broadcast.py b/python/pyspark/broadcast.py
@@ -139,6 +139,29 @@ def __reduce__(self):
         return _from_id, (self._jbroadcast.id(),)
 
 
+class BroadcastPickleRegistry(object):
+    """ Thread-safe registry for broadcast variables that have been pickled
+    """
+
+    def __init__(self, lock):
+        self._registry = set()
+        self._lock = lock
+
+    @property
+    def lock(self):
+        return self._lock
+
+    def add(self, bcast):
+        with self._lock:
+            self._registry.add(bcast)
+
+    def get_and_clear(self):
+        with self._lock:
+            registry_copy = self._registry.copy()
+            self._registry.clear()
+        return registry_copy
+
+
 if __name__ == "__main__":
     import doctest
     (failure_count, test_count) = doctest.testmod()
diff --git a/python/pyspark/context.py b/python/pyspark/context.py
@@ -30,7 +30,7 @@
 
 from pyspark import accumulators
 from pyspark.accumulators import Accumulator
-from pyspark.broadcast import Broadcast
+from pyspark.broadcast import Broadcast, BroadcastPickleRegistry
 from pyspark.conf import SparkConf
 from pyspark.files import SparkFiles
 from pyspark.java_gateway import launch_gateway
@@ -195,7 +195,7 @@ def _do_init(self, master, appName, sparkHome, pyFiles, environment, batchSize,
         # This allows other code to determine which Broadcast instances have
         # been pickled, so it can determine which Java broadcast objects to
         # send.
-        self._pickled_broadcast_vars = set()
+        self._pickled_broadcast_registry = BroadcastPickleRegistry(self._lock)
 
         SparkFiles._sc = self
         root_dir = SparkFiles.getRootDirectory()
@@ -793,7 +793,7 @@ def broadcast(self, value):
         object for reading it in distributed functions. The variable will
         be sent to each cluster only once.
         """
-        return Broadcast(self, value, self._pickled_broadcast_vars)
+        return Broadcast(self, value, self._pickled_broadcast_registry)
 
     def accumulator(self, value, accum_param=None):
         """
diff --git a/python/pyspark/rdd.py b/python/pyspark/rdd.py
@@ -2370,13 +2370,14 @@ def toLocalIterator(self):
 def _prepare_for_python_RDD(sc, command):
     # the serialized command will be compressed by broadcast
     ser = CloudPickleSerializer()
-    pickled_command = ser.dumps(command)
-    if len(pickled_command) > (1 << 20):  # 1M
-        # The broadcast will have same life cycle as created PythonRDD
-        broadcast = sc.broadcast(pickled_command)
-        pickled_command = ser.dumps(broadcast)
-    broadcast_vars = [x._jbroadcast for x in sc._pickled_broadcast_vars]
-    sc._pickled_broadcast_vars.clear()
+    with sc._pickled_broadcast_registry.lock:
+        pickled_command = ser.dumps(command)
+        if len(pickled_command) > (1 << 20):  # 1M
+            # The broadcast will have same life cycle as created PythonRDD
+            broadcast = sc.broadcast(pickled_command)
+            pickled_command = ser.dumps(broadcast)
+        pickled_broadcast_vars = sc._pickled_broadcast_registry.get_and_clear()
+    broadcast_vars = [x._jbroadcast for x in pickled_broadcast_vars]
     return pickled_command, broadcast_vars, sc.environment, sc._python_includes
 
 
