Detect failed fork() calls; improve error logging.

core/src/main/scala/org/apache/spark/api/python/PythonWorkerFactory.scala

@@ -18,7 +18,7 @@
 package org.apache.spark.api.python
 
 import java.io.{DataInputStream, InputStream, OutputStreamWriter}
-import java.net.{InetAddress, ServerSocket, Socket, SocketException}
+import java.net._
 
 import scala.collection.JavaConversions._
 
@@ -64,10 +64,16 @@ private[spark] class PythonWorkerFactory(pythonExec: String, envVars: Map[String
 
       // Attempt to connect, restart and retry once if it fails
       try {
-        new Socket(daemonHost, daemonPort)
+        val socket = new Socket(daemonHost, daemonPort)
+        val launchStatus = new DataInputStream(socket.getInputStream).readInt()
+        if (launchStatus != 0) {
+          logWarning("Python daemon failed to launch worker")
+        }
+        socket
       } catch {
         case exc: SocketException =>
-          logWarning("Python daemon unexpectedly quit, attempting to restart")
+          logWarning("Failed to open socket to Python daemon:", exc)
+          logWarning("Assuming that daemon unexpectedly quit, attempting to restart")
           stopDaemon()
           startDaemon()
           new Socket(daemonHost, daemonPort)

python/pyspark/daemon.py

@@ -67,6 +67,8 @@ def waitSocketClose(sock):
     outfile = os.fdopen(os.dup(sock.fileno()), "a+", 65536)
     exit_code = 0
     try:
+        write_int(0, outfile)  # Acknowledge that the fork was successful
+        outfile.flush()
         worker_main(infile, outfile)
     except SystemExit as exc:
         exit_code = exc.code
@@ -128,16 +130,24 @@ def handle_sigchld(*args):
             if listen_sock in ready_fds:
                 sock, addr = listen_sock.accept()
                 # Launch a worker process
-                if os.fork() == 0:
-                    listen_sock.close()
-                    try:
-                        worker(sock)
-                    except:
-                        traceback.print_exc()
-                        os._exit(1)
+                try:
+                    fork_return_code = os.fork()
+                    if fork_return_code == 0:
+                        listen_sock.close()
+                        try:
+                            worker(sock)
+                        except:
+                            traceback.print_exc()
+                            os._exit(1)
+                        else:
+                            os._exit(0)
                     else:
-                        os._exit(0)
-                else:
+                        sock.close()
+                except OSError as e:
+                    print >> sys.stderr, "Daemon failed to fork PySpark worker: %s" % e
+                    outfile = os.fdopen(os.dup(sock.fileno()), "a+", 65536)
+                    write_int(-1, outfile)  # Signal that the fork failed
+                    outfile.flush()
                     sock.close()
     finally:
         shutdown(1)