Skip to content

Commit

Permalink
[LIVY-995][REPL] JsonParseException is thrown when closing Livy sessi…
Browse files Browse the repository at this point in the history
…on when using python profile (#439)

* [LIVY-995][REPL] JsonParseException is thrown when closing Livy session when using python profile
  • Loading branch information
jianzhenwu authored Jan 25, 2024
1 parent 86fc823 commit 96b1eb6
Show file tree
Hide file tree
Showing 2 changed files with 39 additions and 3 deletions.
16 changes: 13 additions & 3 deletions repl/src/main/scala/org/apache/livy/repl/PythonInterpreter.scala
Original file line number Diff line number Diff line change
Expand Up @@ -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")}")
}
}

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Expand Down

0 comments on commit 96b1eb6

Please sign in to comment.