Github user aarondav commented on a diff in the pull request:
https://github.com/apache/spark/pull/640#discussion_r12266425
--- Diff: core/src/main/scala/org/apache/spark/api/python/PythonRDD.scala
---
@@ -56,122 +56,46 @@ private[spark] class PythonRDD[T: ClassTag](
val env = SparkEnv.get
val worker: Socket = env.createPythonWorker(pythonExec, envVars.toMap)
- // Ensure worker socket is closed on task completion. Closing sockets
is idempotent.
- context.addOnCompleteCallback(() =>
+ // Start a thread to feed the process input from our parent's iterator
+ val writerThread = new WriterThread(env, worker, split, context)
+
+ context.addOnCompleteCallback { () =>
+ writerThread.shutdownOnTaskCompletion()
+
+ // Cleanup the worker socket. This will also cause the Python worker
to exit.
try {
worker.close()
} catch {
case e: Exception => logWarning("Failed to close worker socket", e)
}
- )
-
- @volatile var readerException: Exception = null
- // Start a thread to feed the process input from our parent's iterator
- new Thread("stdin writer for " + pythonExec) {
- override def run() {
+ // The python worker must be destroyed in the event of cancellation
to ensure it unblocks.
+ if (context.interrupted) {
try {
- SparkEnv.set(env)
- val stream = new BufferedOutputStream(worker.getOutputStream,
bufferSize)
- val dataOut = new DataOutputStream(stream)
- // Partition index
- dataOut.writeInt(split.index)
- // sparkFilesDir
- PythonRDD.writeUTF(SparkFiles.getRootDirectory, dataOut)
- // Broadcast variables
- dataOut.writeInt(broadcastVars.length)
- for (broadcast <- broadcastVars) {
- dataOut.writeLong(broadcast.id)
- dataOut.writeInt(broadcast.value.length)
- dataOut.write(broadcast.value)
- }
- // Python includes (*.zip and *.egg files)
- dataOut.writeInt(pythonIncludes.length)
- for (include <- pythonIncludes) {
- PythonRDD.writeUTF(include, dataOut)
- }
- dataOut.flush()
- // Serialized command:
- dataOut.writeInt(command.length)
- dataOut.write(command)
- // Data values
- PythonRDD.writeIteratorToStream(parent.iterator(split, context),
dataOut)
- dataOut.flush()
- worker.shutdownOutput()
+ logWarning("Incomplete task interrupted: Attempting to kill
Python Worker")
+ env.destroyPythonWorker(pythonExec, envVars.toMap)
} catch {
-
- case e: java.io.FileNotFoundException =>
- readerException = e
- Try(worker.shutdownOutput()) // kill Python worker process
-
- case e: IOException =>
- // This can happen for legitimate reasons if the Python code
stops returning data
- // before we are done passing elements through, e.g., for
take(). Just log a message to
- // say it happened (as it could also be hiding a real
IOException from a data source).
- logInfo("stdin writer to Python finished early (may not be an
error)", e)
-
- case e: Exception =>
- // We must avoid throwing exceptions here, because the thread
uncaught exception handler
- // will kill the whole executor (see Executor).
- readerException = e
- Try(worker.shutdownOutput()) // kill Python worker process
+ case e: Exception => logError("Exception when trying to kill
worker", e)
}
}
- }.start()
-
- // Necessary to distinguish between a task that has failed and a task
that is finished
- @volatile var complete: Boolean = false
-
- // It is necessary to have a monitor thread for python workers if the
user cancels with
- // interrupts disabled. In that case we will need to explicitly kill
the worker, otherwise the
- // threads can block indefinitely.
- new Thread(s"Worker Monitor for $pythonExec") {
- override def run() {
- // Kill the worker if it is interrupted or completed
- // When a python task completes, the context is always set to
interupted
- while (!context.interrupted) {
- Thread.sleep(2000)
- }
- if (!complete) {
- try {
- logWarning("Incomplete task interrupted: Attempting to kill
Python Worker")
- env.destroyPythonWorker(pythonExec, envVars.toMap)
- } catch {
- case e: Exception =>
- logError("Exception when trying to kill worker", e)
- }
- }
- }
- }.start()
-
- /*
- * Partial fix for SPARK-1019: Attempts to stop reading the input
stream since
- * other completion callbacks might invalidate the input. Because
interruption
- * is not synchronous this still leaves a potential race where the
interruption is
- * processed only after the stream becomes invalid.
- */
- context.addOnCompleteCallback{ () =>
- complete = true // Indicate that the task has completed successfully
- context.interrupted = true
}
+ writerThread.start()
+
// Return an iterator that read lines from the process's stdout
val stream = new DataInputStream(new
BufferedInputStream(worker.getInputStream, bufferSize))
val stdoutIterator = new Iterator[Array[Byte]] {
def next(): Array[Byte] = {
val obj = _nextObj
if (hasNext) {
- // FIXME: can deadlock if worker is waiting for us to
--- End diff --
@jey I removed this comment because after some effort I could not figure
out how a deadlock could occur since the only "response" mechanism is via the
stdout writer thread which does not interact with this iterator. Please let me
know if this comment is still valid.
---
If your project is set up for it, you can reply to this email and have your
reply appear on GitHub as well. If your project does not have this feature
enabled and wishes so, or if the feature is enabled but not working, please
contact infrastructure at [email protected] or file a JIRA ticket
with INFRA.
---
