Prototype of runtime profiler

Author: HyukjinKwon

core/src/main/scala/org/apache/spark/SparkEnv.scala

@@ -94,7 +94,7 @@ class SparkEnv (
    */
   private case class PythonWorkersKey(
       pythonExec: String, workerModule: String, daemonModule: String, envVars: Map[String, String])
-  private val pythonWorkers = mutable.HashMap[PythonWorkersKey, PythonWorkerFactory]()
+  private[sql] val pythonWorkers = mutable.HashMap[PythonWorkersKey, PythonWorkerFactory]()
 
   // A general, soft-reference map for metadata needed during HadoopRDD split computation
   // (e.g., HadoopFileRDD uses this to cache JobConfs and InputFormats).

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

@@ -38,36 +38,43 @@ import org.apache.spark.internal.config.Python.PYTHON_FACTORY_IDLE_WORKER_MAX_PO
 import org.apache.spark.security.SocketAuthHelper
 import org.apache.spark.util.{RedirectThread, Utils}
 
-case class PythonWorker(channel: SocketChannel) {
-
-  private[this] var selectorOpt: Option[Selector] = None
-  private[this] var selectionKeyOpt: Option[SelectionKey] = None
-
-  def selector: Selector = selectorOpt.orNull
-  def selectionKey: SelectionKey = selectionKeyOpt.orNull
-
-  private def closeSelector(): Unit = {
-    selectionKeyOpt.foreach(_.cancel())
-    selectorOpt.foreach(_.close())
+case class PythonWorker(
+    channel: SocketChannel,
+    extraChannel: Option[SocketChannel] = None) {
+
+  private[this] var selectors: Seq[Selector] = Seq.empty
+  private[this] var selectionKeys: Seq[SelectionKey] = Seq.empty
+
+  private def closeSelectors(): Unit = {
+    selectionKeys.foreach(_.cancel())
+    selectors.foreach(_.close())
+    selectors = Seq.empty
+    selectionKeys = Seq.empty
   }
 
   def refresh(): this.type = synchronized {
-    closeSelector()
-    if (channel.isBlocking) {
-      selectorOpt = None
-      selectionKeyOpt = None
-    } else {
-      val selector = Selector.open()
-      selectorOpt = Some(selector)
-      selectionKeyOpt =
-        Some(channel.register(selector, SelectionKey.OP_READ | SelectionKey.OP_WRITE))
-    }
+    closeSelectors()
+
+    val channels = Seq(Some(channel), extraChannel).flatten
+    val (selList, keyList) = channels.map { ch =>
+      if (ch.isBlocking) {
+        (None, None)
+      } else {
+        val selector = Selector.open()
+        val key = ch.register(selector, SelectionKey.OP_READ | SelectionKey.OP_WRITE)
+        (Some(selector), Some(key))
+      }
+    }.unzip
+
+    selectors = selList.flatten
+    selectionKeys = keyList.flatten
     this
   }
 
   def stop(): Unit = synchronized {
-    closeSelector()
+    closeSelectors()
     Option(channel).foreach(_.close())
+    extraChannel.foreach(_.close())
   }
 }
 
@@ -129,6 +136,10 @@ private[spark] class PythonWorkerFactory(
     envVars.getOrElse("PYTHONPATH", ""),
     sys.env.getOrElse("PYTHONPATH", ""))
 
+  def getAllDaemonWorkers: Seq[(PythonWorker, ProcessHandle)] = self.synchronized {
+    daemonWorkers.filter { case (_, handle) => handle.isAlive}.toSeq
+  }
+
   def create(): (PythonWorker, Option[ProcessHandle]) = {
     if (useDaemon) {
       self.synchronized {
@@ -163,22 +174,36 @@ private[spark] class PythonWorkerFactory(
   private def createThroughDaemon(): (PythonWorker, Option[ProcessHandle]) = {
 
     def createWorker(): (PythonWorker, Option[ProcessHandle]) = {
-      val socketChannel = if (isUnixDomainSock) {
+      val mainChannel = if (isUnixDomainSock) {
         SocketChannel.open(UnixDomainSocketAddress.of(daemonSockPath))
       } else {
         SocketChannel.open(new InetSocketAddress(daemonHost, daemonPort))
       }
+
+      val extraChannel = if (envVars.getOrElse("PYSPARK_RUNTIME_PROFILE", "false").toBoolean) {
+        if (isUnixDomainSock) {
+          Some(SocketChannel.open(UnixDomainSocketAddress.of(daemonSockPath)))
+        } else {
+          Some(SocketChannel.open(new InetSocketAddress(daemonHost, daemonPort)))
+        }
+      } else {
+        None
+      }
+
       // These calls are blocking.
-      val pid = new DataInputStream(Channels.newInputStream(socketChannel)).readInt()
+      val pid = new DataInputStream(Channels.newInputStream(mainChannel)).readInt()
       if (pid < 0) {
         throw new IllegalStateException("Python daemon failed to launch worker with code " + pid)
       }
       val processHandle = ProcessHandle.of(pid).orElseThrow(
         () => new IllegalStateException("Python daemon failed to launch worker.")
       )
-      authHelper.authToServer(socketChannel)
-      socketChannel.configureBlocking(false)
-      val worker = PythonWorker(socketChannel)
+
+      authHelper.authToServer(mainChannel)
+      mainChannel.configureBlocking(false)
+      extraChannel.foreach(_.configureBlocking(false))
+
+      val worker = PythonWorker(mainChannel, extraChannel)
       daemonWorkers.put(worker, processHandle)
       (worker.refresh(), Some(processHandle))
     }
@@ -271,7 +296,7 @@ private[spark] class PythonWorkerFactory(
         if (!blockingMode) {
           socketChannel.configureBlocking(false)
         }
-        val worker = PythonWorker(socketChannel)
+        val worker = PythonWorker(socketChannel, None)
         self.synchronized {
           simpleWorkers.put(worker, workerProcess)
         }

python/pyspark/daemon.py

@@ -47,7 +47,7 @@ def compute_real_exit_code(exit_code):
         return 1
 
 
-def worker(sock, authenticated):
+def worker(sock, sock2, authenticated):
     """
     Called by a worker process after the fork().
     """
@@ -64,6 +64,9 @@ def worker(sock, authenticated):
     buffer_size = int(os.environ.get("SPARK_BUFFER_SIZE", 65536))
     infile = os.fdopen(os.dup(sock.fileno()), "rb", buffer_size)
     outfile = os.fdopen(os.dup(sock.fileno()), "wb", buffer_size)
+    outfile2 = None
+    if sock2 is not None:
+        outfile2 = os.fdopen(os.dup(sock2.fileno()), "wb", buffer_size)
 
     if not authenticated:
         client_secret = UTF8Deserializer().loads(infile)
@@ -74,11 +77,16 @@ def worker(sock, authenticated):
             write_with_length("err".encode("utf-8"), outfile)
             outfile.flush()
             sock.close()
+            if sock2 is not None:
+                sock2.close()
             return 1
 
     exit_code = 0
     try:
-        worker_main(infile, outfile)
+        if sock2 is not None:
+            worker_main(infile, (outfile, outfile2))
+        else:
+            worker_main(infile, outfile)
     except SystemExit as exc:
         exit_code = compute_real_exit_code(exc.code)
     finally:
@@ -94,6 +102,7 @@ def manager():
     os.setpgid(0, 0)
 
     is_unix_domain_sock = os.environ.get("PYTHON_UNIX_DOMAIN_ENABLED", "false").lower() == "true"
+    is_python_runtime_profile = os.environ.get("PYSPARK_RUNTIME_PROFILE", "false").lower() == "true"
     socket_path = None
 
     # Create a listening socket on the loopback interface
@@ -173,6 +182,15 @@ def handle_sigterm(*args):
                         continue
                     raise
 
+                sock2 = None
+                if is_python_runtime_profile:
+                    try:
+                        sock2, _ = listen_sock.accept()
+                    except OSError as e:
+                        if e.errno == EINTR:
+                            continue
+                        raise
+
                 # Launch a worker process
                 try:
                     pid = os.fork()
@@ -186,6 +204,13 @@ def handle_sigterm(*args):
                         outfile.flush()
                         outfile.close()
                         sock.close()
+
+                        if sock2 is not None:
+                            outfile = sock2.makefile(mode="wb")
+                            write_int(e.errno, outfile)  # Signal that the fork failed
+                            outfile.flush()
+                            outfile.close()
+                            sock2.close()
                         continue
 
                 if pid == 0:
@@ -217,14 +242,16 @@ def handle_sigterm(*args):
                             or False
                         )
                         while True:
-                            code = worker(sock, authenticated)
+                            code = worker(sock, sock2, authenticated)
                             if code == 0:
                                 authenticated = True
                             if not reuse or code:
                                 # wait for closing
                                 try:
                                     while sock.recv(1024):
                                         pass
+                                    while sock2 is not None and sock2.recv(1024):
+                                        pass
                                 except Exception:
                                     pass
                                 break

python/pyspark/worker.py

@@ -18,6 +18,8 @@
 """
 Worker that receives input from Piped RDD.
 """
+import pickle
+import threading
 import itertools
 import os
 import sys
@@ -45,6 +47,7 @@
     read_bool,
     write_long,
     read_int,
+    write_with_length,
     SpecialLengths,
     CPickleSerializer,
     BatchedSerializer,
@@ -3167,7 +3170,27 @@ def func(_, it):
     return func, None, ser, ser
 
 
+def write_profile(outfile):
+    import yappi
+
+    while True:
+        stats = []
+        for thread in yappi.get_thread_stats():
+            data = list(yappi.get_func_stats(ctx_id=thread.id))
+            stats.extend([{str(k): str(v) for k, v in d.items()} for d in data])
+        pickled = pickle.dumps(stats)
+        write_with_length(pickled, outfile)
+        time.sleep(1)
+
+
 def main(infile, outfile):
+    if isinstance(outfile, tuple):
+        import yappi
+
+        outfile, outfile2 = outfile
+        yappi.start()
+        threading.Thread(target=write_profile, args=(outfile2,), daemon=True).start()
+
     faulthandler_log_path = os.environ.get("PYTHON_FAULTHANDLER_DIR", None)
     tracebackDumpIntervalSeconds = os.environ.get("PYTHON_TRACEBACK_DUMP_INTERVAL_SECONDS", None)
     try:

sql/core/src/main/resources/META-INF/services/org.apache.spark.sql.sources.DataSourceRegister

@@ -26,6 +26,7 @@ org.apache.spark.sql.execution.datasources.xml.XmlFileFormat
 org.apache.spark.sql.execution.streaming.ConsoleSinkProvider
 org.apache.spark.sql.execution.streaming.sources.RateStreamProvider
 org.apache.spark.sql.execution.streaming.sources.TextSocketSourceProvider
+org.apache.spark.sql.execution.streaming.sources.PythonProfileSourceProvider
 org.apache.spark.sql.execution.datasources.binaryfile.BinaryFileFormat
 org.apache.spark.sql.execution.streaming.sources.RatePerMicroBatchProvider
 org.apache.spark.sql.execution.datasources.v2.state.metadata.StateMetadataSource