瀏覽代碼

[livy] Add support for python to livy-server

Erick Tryzelaar 10 年之前
父節點
當前提交
28068798b8

+ 1 - 1
apps/spark/java/bin/spark-shell → apps/spark/java/bin/livy-repl

@@ -23,4 +23,4 @@ export LIVY_HOME=$(cd $(dirname $0)/.. && pwd)
 
 exec java \
 	-cp "$LIVY_HOME/livy-repl/target/lib/*:$LIVY_HOME/livy-repl/target/livy-repl-3.7.0-SNAPSHOT.jar" \
-	com.cloudera.hue.livy.repl.Main -usejavacp "$@"
+	com.cloudera.hue.livy.repl.Main "$@"

+ 180 - 0
apps/spark/java/livy-repl/src/main/python/fake_shell.py

@@ -0,0 +1,180 @@
+import cStringIO
+import json
+import logging
+import sys
+import traceback
+
+logging.basicConfig()
+logger = logging.getLogger('fake_shell')
+
+sys_stdin = sys.stdin
+sys_stdout = sys.stdout
+sys_stderr = sys.stderr
+
+fake_stdin = cStringIO.StringIO()
+fake_stdout = cStringIO.StringIO()
+fake_stderr = cStringIO.StringIO()
+
+sys.stdin = fake_stdin
+sys.stdout = fake_stdout
+sys.stderr = fake_stderr
+
+global_dict = {}
+
+execution_count = 0
+
+def execute_request(msg):
+    global execution_count
+
+    try:
+        code = msg['code']
+    except KeyError:
+        logger.error('missing code', exc_info=True)
+        return
+
+    execution_count += 1
+
+    try:
+        code = compile(code, '<stdin>', 'single')
+        exec code in global_dict
+    except:
+        exc_type, exc_value, tb = sys.exc_info()
+        return {
+            'msg_type': 'execute_reply',
+            'execution_count': execution_count - 1,
+            'content': {
+                'status': 'error',
+                'ename': exc_type.__name__,
+                'evalue': str(exc_value),
+                'traceback': traceback.extract_tb(tb),
+            }
+        }
+
+    stdout = fake_stdout.getvalue()
+    stderr = fake_stderr.getvalue()
+
+    output = ''
+
+    if stdout:
+        output += stdout
+
+    if stderr:
+        output += stderr
+
+    return {
+        'msg_type': 'execute_result',
+        'execution_count': execution_count,
+        'content': {
+            'status': 'ok',
+            'execution_count': execution_count - 1,
+            'data': {
+                'text/plain': output,
+            },
+        }
+    }
+
+def inspect_value(name):
+    try:
+        value = global_dict[name]
+    except KeyError:
+        return {
+            'msg_type': 'inspect_reply',
+            'execution_count': execution_count - 1,
+            'content': {
+                'status': 'error',
+                'ename': 'KeyError',
+                'evalue': 'unknown variable %s' % name,
+            }
+        }
+
+    return {
+        'msg_type': 'inspect_result',
+        'content': {
+            'data': {
+                'application/json': value,
+            },
+        }
+    }
+
+
+inspect_router = {
+    '%inspect': inspect_value,
+}
+
+def inspect_request(msg):
+    try:
+        code = msg['code']
+    except KeyError:
+        logger.error('missing code', exc_info=True)
+        return
+
+    try:
+        inspect_magic, code = code.split(' ', 1)
+    except ValueError:
+        logger.error('invalid magic', exc_info=True)
+        return
+
+    try:
+        handler = inspect_router[inspect_magic]
+    except KeyError:
+        logger.error('unknown magic', exc_info=True)
+        return
+
+    return handler(code)
+
+
+msg_type_router = {
+    'execute_request': execute_request,
+    'inspect_request': inspect_request,
+}
+
+try:
+    while True:
+        fake_stdout.truncate(0)
+
+        line = sys_stdin.readline()
+
+        if line == '':
+            break
+        elif line == '\n':
+            continue
+
+        try:
+            msg = json.loads(line)
+        except ValueError:
+            logger.error('failed to parse message', exc_info=True)
+            continue
+
+        try:
+            msg_type = msg['msg_type']
+        except KeyError:
+            logger.error('missing message type')
+            continue
+
+        try:
+            handler = msg_type_router[msg_type]
+        except KeyError:
+            logger.error('unknown message type: %s', msg_type)
+            continue
+
+        response = handler(msg)
+        if response is not None:
+            try:
+                response = json.dumps(response)
+            except ValueError:
+                response = json.dumps({
+                    'msg_type': 'inspect_reply',
+                    'execution_count': execution_count - 1,
+                    'content': {
+                        'status': 'error',
+                        'ename': 'ValueError',
+                        'evalue': 'cannot json-ify %s' % name,
+                    }
+                })
+
+            print >> sys_stdout, response
+            sys_stdout.flush()
+finally:
+    sys.stdin = sys_stdin
+    sys.stdout = sys_stdout
+    sys.stderr = sys_stderr

+ 29 - 2
apps/spark/java/livy-repl/src/main/scala/com/cloudera/hue/livy/repl/Main.scala

@@ -2,19 +2,41 @@ package com.cloudera.hue.livy.repl
 
 import javax.servlet.ServletContext
 
-import com.cloudera.hue.livy.repl.spark.SparkSession
+import com.cloudera.hue.livy.repl.python.PythonSession
+import com.cloudera.hue.livy.repl.scala.ScalaSession
 import com.cloudera.hue.livy.{Logging, WebServer}
 import org.scalatra.LifeCycle
 import org.scalatra.servlet.ScalatraListener
 
 object Main extends Logging {
+
+  val SESSION_KIND = "livy-repl.session.kind"
+  val PYTHON_SESSION = "python"
+  val SCALA_SESSION = "scala"
+
   def main(args: Array[String]): Unit = {
     val port = sys.env.getOrElse("PORT", "8999").toInt
+
+    if (args.length != 1) {
+      println("Must specify either `python` or `scala` for the session kind")
+      sys.exit(1)
+    }
+
+    val session_kind = args(0)
+
+    session_kind match {
+      case PYTHON_SESSION | SCALA_SESSION =>
+      case _ =>
+        println("Unknown session kind: " + session_kind)
+        sys.exit(1)
+    }
+
     val server = new WebServer(port)
 
     server.context.setResourceBase("src/main/com/cloudera/hue/livy/repl")
     server.context.setInitParameter(ScalatraListener.LifeCycleKey, classOf[ScalatraBootstrap].getCanonicalName)
     server.context.addEventListener(new ScalatraListener)
+    server.context.setInitParameter(SESSION_KIND, session_kind)
 
     server.start()
     println("Starting livy-repl on port %s" format server.port)
@@ -26,9 +48,14 @@ object Main extends Logging {
 
 class ScalatraBootstrap extends LifeCycle {
 
-  val session = new SparkSession()
+  var session: Session = null
 
   override def init(context: ServletContext): Unit = {
+    val session = context.getInitParameter(Main.SESSION_KIND) match {
+      case Main.PYTHON_SESSION => PythonSession.create()
+      case Main.SCALA_SESSION => ScalaSession.create()
+    }
+
     context.mount(new WebApp(session), "/*")
   }
 

+ 1 - 1
apps/spark/java/livy-repl/src/main/scala/com/cloudera/hue/livy/repl/Session.scala

@@ -2,7 +2,7 @@ package com.cloudera.hue.livy.repl
 
 import com.cloudera.hue.livy.ExecuteResponse
 
-import scala.concurrent.Future
+import _root_.scala.concurrent.Future
 
 trait Session {
   def statements: Seq[ExecuteResponse]

+ 5 - 5
apps/spark/java/livy-repl/src/main/scala/com/cloudera/hue/livy/repl/WebApp.scala

@@ -1,20 +1,20 @@
 package com.cloudera.hue.livy.repl
 
 import _root_.akka.util.Timeout
-import com.cloudera.hue.livy.{Logging, ExecuteRequest}
+import com.cloudera.hue.livy.{ExecuteRequest, Logging}
 import com.fasterxml.jackson.core.JsonParseException
 import org.json4s.{DefaultFormats, Formats, MappingException}
-import org.scalatra.json.JacksonJsonSupport
 import org.scalatra._
+import org.scalatra.json.JacksonJsonSupport
 
-import scala.concurrent.{ExecutionContext, ExecutionContextExecutor, Future}
+import _root_.scala.concurrent.{Future, ExecutionContext}
 
 object WebApp extends Logging {}
 
 class WebApp(session: Session) extends ScalatraServlet with FutureSupport with JacksonJsonSupport {
 
-  override protected implicit def executor: ExecutionContextExecutor = ExecutionContext.global
-  override protected implicit val jsonFormats: Formats = DefaultFormats
+  override protected implicit def executor = ExecutionContext.global
+  override protected implicit val jsonFormats = DefaultFormats
 
   protected implicit def defaultTimeout: Timeout = Timeout(10)
 

+ 164 - 0
apps/spark/java/livy-repl/src/main/scala/com/cloudera/hue/livy/repl/python/PythonSession.scala

@@ -0,0 +1,164 @@
+package com.cloudera.hue.livy.repl.python
+
+import java.io._
+
+import com.cloudera.hue.livy.ExecuteResponse
+import com.cloudera.hue.livy.repl.Session
+import org.json4s.DefaultFormats
+import org.json4s.JsonAST._
+import org.json4s.JsonDSL._
+import org.json4s.jackson.JsonMethods._
+
+import scala.collection.mutable.ArrayBuffer
+import scala.concurrent.{ExecutionContext, Future}
+
+object PythonSession {
+  val LIVY_HOME = System.getenv("LIVY_HOME")
+  val FAKE_SHELL = LIVY_HOME + "/livy-repl/src/main/python/fake_shell.py"
+
+  def create(): Session = {
+    val pb = new ProcessBuilder("python", FAKE_SHELL)
+    val process = pb.start()
+    val in = process.getInputStream
+    val out = process.getOutputStream
+
+    new PythonSession(process, in, out)
+  }
+
+  // Java unfortunately wraps the input stream in a buffer, so we need to hack around it so we can read the output
+  // without blocking.
+  private def unwrapInputStream(inputStream: InputStream) = {
+    var filteredInputStream = inputStream
+
+    while (filteredInputStream.isInstanceOf[FilterInputStream]) {
+      val field = classOf[FilterInputStream].getDeclaredField("in")
+      field.setAccessible(true)
+      filteredInputStream = field.get(filteredInputStream).asInstanceOf[InputStream]
+    }
+
+    filteredInputStream
+  }
+
+  // Java unfortunately wraps the output stream in a buffer, so we need to hack around it so we can read the output
+  // without blocking.
+  private def unwrapOutputStream(outputStream: OutputStream) = {
+    var filteredOutputStream = outputStream
+
+    while (filteredOutputStream.isInstanceOf[FilterOutputStream]) {
+      val field = classOf[FilterOutputStream].getDeclaredField("out")
+      field.setAccessible(true)
+      filteredOutputStream = field.get(filteredOutputStream).asInstanceOf[OutputStream]
+    }
+
+    filteredOutputStream
+  }
+}
+
+private class PythonSession(process: Process, in: InputStream, out: OutputStream) extends Session {
+  private implicit def executor: ExecutionContext = ExecutionContext.global
+
+  implicit val formats = DefaultFormats
+
+  private[this] val stdin = new PrintWriter(out)
+  private[this] val stdout = new BufferedReader(new InputStreamReader(in), 1)
+
+  private[this] var executedStatements = 0
+  private[this] var _statements = ArrayBuffer[ExecuteResponse]()
+
+  override def statements: Seq[ExecuteResponse] = _statements
+
+  override def execute(command: String): Future[ExecuteResponse] = {
+    val request = Map(
+      "msg_type" -> "execute_request",
+      "code" -> command
+    )
+
+    stdin.println(compact(render(request)))
+    stdin.flush()
+
+    val line = stdout.readLine()
+
+    val response = parse(line)
+
+    val content = response \ "content"
+    val status = (content\ "status").extract[String]
+    val executionCount = (content \ "execution_count").extract[Int]
+
+    val executeResponse = status match {
+      case "ok" =>
+        val output = (content \ "data" \ "text/plain").extract[String]
+        ExecuteResponse(executionCount, Seq(command), Seq(output))
+      case "error" =>
+        val ename = (content \ "ename").extract[String]
+        val evalue = (content \ "evalue").extract[String]
+        val traceback = (content \ "traceback").extract[Seq[String]]
+
+        val output = traceback :+ ("%s: %s" format(ename, evalue))
+
+        ExecuteResponse(executionCount, Seq(command), output)
+    }
+
+    Future.successful(executeResponse)
+
+    /*
+    val response = ExecuteResponse(executedStatements - 1, Seq(command), output)
+    _statements += response
+
+    executedStatements += 1
+
+    Future.successful(response)
+    */
+  }
+
+  override def statement(id: Int): Option[ExecuteResponse] = {
+    if (id < _statements.length) {
+      Some(_statements(id))
+    } else {
+      None
+    }
+  }
+
+  override def close(): Unit = {
+    process.getInputStream.close()
+    process.getOutputStream.close()
+    process.destroy()
+    /*
+    if (!process.waitFor(10l, TimeUnit.SECONDS)) {
+      process.destroyForcibly()
+      process.waitFor()
+    }
+    */
+  }
+
+
+  /*
+  private def readLines(): Seq[String] = {
+    var sb = new StringBuilder
+    var output = new ArrayBuffer[String]()
+
+    @tailrec
+    def aux(): Unit = {
+      stdout.
+
+
+      val line = stdout.readLine()
+      if (line != null && !line.startsWith(">>> ") && !line.startsWith("... ")) {
+        output += line
+        aux()
+      }
+    }
+
+    aux()
+
+    output
+
+    /*
+    val output = stdout.takeWhile({
+      case line: String =>
+        println(line)
+        line.startsWith(">>> ") || line.startsWith("... ")
+    }).toSeq
+    */
+  }
+  */
+}

+ 1 - 1
apps/spark/java/livy-repl/src/main/scala/com/cloudera/hue/livy/repl/spark/ILoop.scala → apps/spark/java/livy-repl/src/main/scala/com/cloudera/hue/livy/repl/scala/ILoop.scala

@@ -1,4 +1,4 @@
-package com.cloudera.hue.livy.repl.spark
+package com.cloudera.hue.livy.repl.scala
 
 import java.io._
 import java.util.concurrent.BlockingQueue

+ 6 - 2
apps/spark/java/livy-repl/src/main/scala/com/cloudera/hue/livy/repl/spark/SparkSession.scala → apps/spark/java/livy-repl/src/main/scala/com/cloudera/hue/livy/repl/scala/ScalaSession.scala

@@ -1,4 +1,4 @@
-package com.cloudera.hue.livy.repl.spark
+package com.cloudera.hue.livy.repl.scala
 
 import java.util.concurrent.SynchronousQueue
 
@@ -9,7 +9,11 @@ import scala.collection.mutable
 import scala.concurrent.duration.Duration
 import scala.concurrent.{Await, ExecutionContext, Future, Promise}
 
-class SparkSession extends Session {
+object ScalaSession {
+  def create(): Session = new ScalaSession()
+}
+
+private class ScalaSession extends Session {
   private implicit def executor: ExecutionContext = ExecutionContext.global
 
   private[this] val inQueue = new SynchronousQueue[ILoop.Request]

+ 1 - 1
apps/spark/java/livy-server/src/main/scala/com/cloudera/hue/livy/server/Main.scala

@@ -8,7 +8,7 @@ import org.scalatra.servlet.ScalatraListener
 
 object Main {
 
-  val SESSION_KIND = "livy.session.kind"
+  val SESSION_KIND = "livy-server.session.kind"
   val PROCESS_SESSION = "process"
   val YARN_SESSION = "yarn"
 

+ 6 - 6
apps/spark/java/livy-server/src/main/scala/com/cloudera/hue/livy/server/SessionFactory.scala

@@ -2,14 +2,14 @@ package com.cloudera.hue.livy.server
 
 import java.util.UUID
 
-import com.cloudera.hue.livy.server.sessions.{Session, SparkYarnSession, SparkProcessSession}
+import com.cloudera.hue.livy.server.sessions._
 import com.cloudera.hue.livy.yarn.Client
 import org.apache.hadoop.yarn.conf.YarnConfiguration
 
 import scala.concurrent.{ExecutionContext, Future}
 
 trait SessionFactory {
-  def createSparkSession(): Future[Session]
+  def createSession(lang: String): Future[Session]
 
   def close(): Unit = {}
 }
@@ -18,10 +18,10 @@ class ProcessSessionFactory extends SessionFactory {
 
   implicit def executor: ExecutionContext = ExecutionContext.global
 
-  override def createSparkSession(): Future[Session] = {
+  override def createSession(lang: String): Future[Session] = {
     Future {
       val id = UUID.randomUUID().toString
-      SparkProcessSession.create(id)
+      ProcessSession.create(id, lang)
     }
   }
 }
@@ -33,9 +33,9 @@ class YarnSessionFactory extends SessionFactory {
 
   val client = new Client(yarnConf)
 
-  override def createSparkSession(): Future[Session] = {
+  override def createSession(lang: String): Future[Session] = {
     val id = UUID.randomUUID().toString
-    SparkYarnSession.create(client, id)
+    YarnSession.create(client, id, lang)
   }
 
   override def close(): Unit = {

+ 2 - 2
apps/spark/java/livy-server/src/main/scala/com/cloudera/hue/livy/server/SessionManager.scala

@@ -32,8 +32,8 @@ class SessionManager(factory: SessionFactory) extends Logging {
     sessions.keys
   }
 
-  def createSparkSession(): Future[Session] = {
-    val session = factory.createSparkSession()
+  def createSession(lang: String): Future[Session] = {
+    val session = factory.createSession(lang)
 
     session.map({ case(session: Session) =>
       info("created session %s" format session.id)

+ 2 - 2
apps/spark/java/livy-server/src/main/scala/com/cloudera/hue/livy/server/WebApp.scala

@@ -46,7 +46,8 @@ class WebApp(sessionManager: SessionManager)
     val createSessionRequest = parsedBody.extract[CreateSessionRequest]
 
     val sessionFuture = createSessionRequest.lang match {
-      case "scala" => sessionManager.createSparkSession()
+      case "scala" => sessionManager.createSession(createSessionRequest.lang)
+      case "python" => sessionManager.createSession(createSessionRequest.lang)
       case lang => halt(400, "unsupported language: " + lang)
     }
 
@@ -87,7 +88,6 @@ class WebApp(sessionManager: SessionManager)
       _ <- sessionManager.delete(params("sessionId"))
     } yield Accepted()
 
-    // FIXME: this is silently eating exceptions.
     new AsyncResult() { val is = for { _ <- future } yield NoContent }
   }
 

+ 10 - 10
apps/spark/java/livy-server/src/main/scala/com/cloudera/hue/livy/server/sessions/SparkProcessSession.scala → apps/spark/java/livy-server/src/main/scala/com/cloudera/hue/livy/server/sessions/ProcessSession.scala

@@ -8,17 +8,17 @@ import scala.annotation.tailrec
 import scala.concurrent.Future
 import scala.io.Source
 
-object SparkProcessSession extends Logging {
+object ProcessSession extends Logging {
   val LIVY_HOME = System.getenv("LIVY_HOME")
-  val SPARK_SHELL = LIVY_HOME + "/bin/spark-shell"
+  val LIVY_REPL = LIVY_HOME + "/bin/livy-repl"
 
-  def create(id: String): Session = {
-    val (process, port) = startProcess()
-    new SparkProcessSession(id, process, port)
+  def create(id: String, lang: String): Session = {
+    val (process, port) = startProcess(lang)
+    new ProcessSession(id, process, port)
   }
 
   // Loop until we've started a process with a valid port.
-  private def startProcess(): (Process, Int) = {
+  private def startProcess(lang: String): (Process, Int) = {
     val regex = """Starting livy-repl on port (\d+)""".r
 
     @tailrec
@@ -36,8 +36,8 @@ object SparkProcessSession extends Logging {
       }
     }
 
-    def startProcess(): (Process, Int) = {
-      val pb = new ProcessBuilder(SPARK_SHELL)
+    def startProcess(lang: String): (Process, Int) = {
+      val pb = new ProcessBuilder(LIVY_REPL, lang)
       pb.environment().put("PORT", "0")
       pb.redirectError(Redirect.INHERIT)
       val process = pb.start()
@@ -58,11 +58,11 @@ object SparkProcessSession extends Logging {
       }
     }
 
-    startProcess()
+    startProcess(lang)
   }
 }
 
-private class SparkProcessSession(id: String, process: Process, port: Int) extends SparkWebSession(id, "localhost", port) {
+private class ProcessSession(id: String, process: Process, port: Int) extends WebSession(id, "localhost", port) {
 
   override def stop(): Future[Unit] = {
     super.stop() andThen { case r =>

+ 1 - 1
apps/spark/java/livy-server/src/main/scala/com/cloudera/hue/livy/server/sessions/SparkWebSession.scala → apps/spark/java/livy-server/src/main/scala/com/cloudera/hue/livy/server/sessions/WebSession.scala

@@ -10,7 +10,7 @@ import scala.annotation.tailrec
 import scala.collection.mutable.ArrayBuffer
 import scala.concurrent.{Future, _}
 
-abstract class SparkWebSession(val id: String, hostname: String, port: Int) extends Session with Logging {
+abstract class WebSession(val id: String, hostname: String, port: Int) extends Session with Logging {
 
   protected implicit def executor: ExecutionContextExecutor = ExecutionContext.global
   protected implicit def jsonFormats: Formats = DefaultFormats

+ 7 - 6
apps/spark/java/livy-server/src/main/scala/com/cloudera/hue/livy/server/sessions/SparkYarnSession.scala → apps/spark/java/livy-server/src/main/scala/com/cloudera/hue/livy/server/sessions/YarnSession.scala

@@ -6,18 +6,19 @@ import org.apache.hadoop.yarn.api.ApplicationConstants
 
 import scala.concurrent.{ExecutionContext, ExecutionContextExecutor, Future, TimeoutException}
 
-object SparkYarnSession {
+object YarnSession {
   private val LIVY_YARN_PACKAGE = System.getenv("LIVY_YARN_PACKAGE")
 
   protected implicit def executor: ExecutionContextExecutor = ExecutionContext.global
 
-  def create(client: Client, id: String): Future[Session] = {
+  def create(client: Client, id: String, lang: String): Future[Session] = {
     val packagePath = new Path(LIVY_YARN_PACKAGE)
 
     val job = client.submitApplication(
       packagePath,
       List(
-        "__package/bin/run-am.sh 1>%s/stdout 2>%s/stderr" format (
+        "__package/bin/run-am.sh %s 1>%s/stdout 2>%s/stderr" format (
+          lang,
           ApplicationConstants.LOG_DIR_EXPANSION_VAR,
           ApplicationConstants.LOG_DIR_EXPANSION_VAR
           )
@@ -31,7 +32,7 @@ object SparkYarnSession {
 
       x match {
         case Some((hostname, port)) =>
-          new SparkYarnSession(id, job, hostname, port)
+          new YarnSession(id, job, hostname, port)
         case None =>
           throw new TimeoutException()
       }
@@ -39,8 +40,8 @@ object SparkYarnSession {
   }
 }
 
-private class SparkYarnSession(id: String, job: Job, hostname: String, port: Int)
-  extends SparkWebSession(id, hostname, port) {
+private class YarnSession(id: String, job: Job, hostname: String, port: Int)
+  extends WebSession(id, hostname, port) {
 
   override def stop(): Future[Unit] = {
     super.stop() andThen { case r =>