diff --git a/repl/src/main/scala/org/apache/livy/repl/PythonInterpreter.scala b/repl/src/main/scala/org/apache/livy/repl/PythonInterpreter.scala index 36c09deb6..58b7147a5 100644 --- a/repl/src/main/scala/org/apache/livy/repl/PythonInterpreter.scala +++ b/repl/src/main/scala/org/apache/livy/repl/PythonInterpreter.scala @@ -269,11 +269,21 @@ private class PythonInterpreter( } override protected def sendShutdownRequest(): Unit = { - sendRequest(Map( + stdin.println(write(Map( "msg_type" -> "shutdown_request", "content" -> () - )).foreach { case rep => - warn(f"process failed to shut down while returning $rep") + ))) + stdin.flush() + + // Pyspark prints profile info to stdout when enabling spark.python.profile. see SPARK-37443 + var lines = Seq[String]() + var line = stdout.readLine() + while(line != null) { + lines :+= line + line = stdout.readLine() + } + if (lines.nonEmpty) { + warn(f"python process shut down while returning ${lines.mkString("\n")}") } } diff --git a/repl/src/test/scala/org/apache/livy/repl/PythonInterpreterSpec.scala b/repl/src/test/scala/org/apache/livy/repl/PythonInterpreterSpec.scala index 4a78c61ff..52a429180 100644 --- a/repl/src/test/scala/org/apache/livy/repl/PythonInterpreterSpec.scala +++ b/repl/src/test/scala/org/apache/livy/repl/PythonInterpreterSpec.scala @@ -249,6 +249,32 @@ abstract class PythonBaseInterpreterSpec extends BaseInterpreterSpec { ) )) } + + it should "work when interpreter exit with json stdout" in { + noException should be thrownBy { + withInterpreter { intp => + val response = intp.execute( + """import atexit, sys + |atexit.register(sys.stdout.write, '{}') + |""".stripMargin + ) + response shouldBe a[Interpreter.ExecuteSuccess] + } + } + } + + it should "work when interpreter exit with non-json stdout" in { + noException should be thrownBy { + withInterpreter { intp => + val response = intp.execute( + """import atexit, sys + |atexit.register(sys.stdout.write, 'line1\nline2') + |""".stripMargin + ) + response shouldBe a[Interpreter.ExecuteSuccess] + } + } + } } class Python2InterpreterSpec extends PythonBaseInterpreterSpec {