From ec31e68d5df259e6df001529235d8c906ff02a6f Mon Sep 17 00:00:00 2001
From: root <root@ip-10-165-19-218.ec2.internal>
Date: Mon, 1 Jul 2013 06:20:14 +0000
Subject: Fixed PySpark perf regression by not using socket.makefile(), and
 improved debuggability by letting "print" statements show up in the
 executor's stderr

Conflicts:
	core/src/main/scala/spark/api/python/PythonRDD.scala
---
 python/pyspark/daemon.py | 42 ++++++++++++++++++++++++------------------
 1 file changed, 24 insertions(+), 18 deletions(-)

(limited to 'python/pyspark/daemon.py')

diff --git a/python/pyspark/daemon.py b/python/pyspark/daemon.py
index 78a2da1e18..78c9457b84 100644
--- a/python/pyspark/daemon.py
+++ b/python/pyspark/daemon.py
@@ -1,10 +1,13 @@
 import os
+import signal
+import socket
 import sys
+import traceback
 import multiprocessing
 from ctypes import c_bool
 from errno import EINTR, ECHILD
-from socket import socket, AF_INET, SOCK_STREAM, SOMAXCONN
-from signal import signal, SIGHUP, SIGTERM, SIGCHLD, SIG_DFL, SIG_IGN
+from socket import AF_INET, SOCK_STREAM, SOMAXCONN
+from signal import SIGHUP, SIGTERM, SIGCHLD, SIG_DFL, SIG_IGN
 from pyspark.worker import main as worker_main
 from pyspark.serializers import write_int
 
@@ -33,11 +36,12 @@ def compute_real_exit_code(exit_code):
 def worker(listen_sock):
     # Redirect stdout to stderr
     os.dup2(2, 1)
+    sys.stdout = sys.stderr   # The sys.stdout object is different from file descriptor 1
 
     # Manager sends SIGHUP to request termination of workers in the pool
     def handle_sighup(*args):
         assert should_exit()
-    signal(SIGHUP, handle_sighup)
+    signal.signal(SIGHUP, handle_sighup)
 
     # Cleanup zombie children
     def handle_sigchld(*args):
@@ -51,7 +55,7 @@ def worker(listen_sock):
                 handle_sigchld()
             elif err.errno != ECHILD:
                 raise
-    signal(SIGCHLD, handle_sigchld)
+    signal.signal(SIGCHLD, handle_sigchld)
 
     # Handle clients
     while not should_exit():
@@ -70,19 +74,22 @@ def worker(listen_sock):
             # never receives SIGCHLD unless a worker crashes.
             if os.fork() == 0:
                 # Leave the worker pool
-                signal(SIGHUP, SIG_DFL)
+                signal.signal(SIGHUP, SIG_DFL)
                 listen_sock.close()
-                # Handle the client then exit
-                sockfile = sock.makefile()
+                # Read the socket using fdopen instead of socket.makefile() because the latter
+                # seems to be very slow; note that we need to dup() the file descriptor because
+                # otherwise writes also cause a seek that makes us miss data on the read side.
+                infile = os.fdopen(os.dup(sock.fileno()), "a+", 65536)
+                outfile = os.fdopen(os.dup(sock.fileno()), "a+", 65536)
                 exit_code = 0
                 try:
-                  worker_main(sockfile, sockfile)
+                    worker_main(infile, outfile)
                 except SystemExit as exc:
-                  exit_code = exc.code
+                    exit_code = exc.code
                 finally:
-                  sockfile.close()
-                  sock.close()
-                  os._exit(compute_real_exit_code(exit_code))
+                    outfile.flush()
+                    sock.close()
+                    os._exit(compute_real_exit_code(exit_code))
             else:
                 sock.close()
 
@@ -92,7 +99,6 @@ def launch_worker(listen_sock):
         try:
             worker(listen_sock)
         except Exception as err:
-            import traceback
             traceback.print_exc()
             os._exit(1)
         else:
@@ -105,7 +111,7 @@ def manager():
     os.setpgid(0, 0)
 
     # Create a listening socket on the AF_INET loopback interface
-    listen_sock = socket(AF_INET, SOCK_STREAM)
+    listen_sock = socket.socket(AF_INET, SOCK_STREAM)
     listen_sock.bind(('127.0.0.1', 0))
     listen_sock.listen(max(1024, 2 * POOLSIZE, SOMAXCONN))
     listen_host, listen_port = listen_sock.getsockname()
@@ -121,8 +127,8 @@ def manager():
         exit_flag.value = True
 
     # Gracefully exit on SIGTERM, don't die on SIGHUP
-    signal(SIGTERM, lambda signum, frame: shutdown())
-    signal(SIGHUP, SIG_IGN)
+    signal.signal(SIGTERM, lambda signum, frame: shutdown())
+    signal.signal(SIGHUP, SIG_IGN)
 
     # Cleanup zombie children
     def handle_sigchld(*args):
@@ -133,7 +139,7 @@ def manager():
         except EnvironmentError as err:
             if err.errno not in (ECHILD, EINTR):
                 raise
-    signal(SIGCHLD, handle_sigchld)
+    signal.signal(SIGCHLD, handle_sigchld)
 
     # Initialization complete
     sys.stdout.close()
@@ -148,7 +154,7 @@ def manager():
                     shutdown()
                     raise
     finally:
-        signal(SIGTERM, SIG_DFL)
+        signal.signal(SIGTERM, SIG_DFL)
         exit_flag.value = True
         # Send SIGHUP to notify workers of shutdown
         os.kill(0, SIGHUP)
-- 
cgit v1.2.3