twosigma · scottdraves · May 16, 2018 · May 15, 2018 · May 15, 2018 · May 16, 2018
diff --git a/beakerx/beakerx/commands.py b/beakerx/beakerx/commands.py
@@ -57,10 +57,11 @@ def py4j_server_subparser(subparser):
     py4j_server_parser.set_defaults(func=start_py4j_server)
     py4j_server_parser.add_argument("--port")
     py4j_server_parser.add_argument("--pyport")
+    py4j_server_parser.add_argument("--kernel")
 
 
 def start_py4j_server(args):
-    Py4JServer(args.port, args.pyport)
+    Py4JServer(args.port, args.pyport, args.kernel)
 
 
 def run_jupyter(jupyter_commands):

diff --git a/beakerx/beakerx_magics/__init__.py b/beakerx/beakerx_magics/__init__.py
@@ -19,4 +19,4 @@
 from .kotlin_magic import *
 from .scala_magic import *
 from .sql_magic import *
-from .python_magic import *
+from .jvm_kernel_magic import *
diff --git a/beakerx/beakerx_magics/python_magic.py → beakerx/beakerx_magics/jvm_kernel_magic.py b/beakerx/beakerx_magics/python_magic.py → beakerx/beakerx_magics/jvm_kernel_magic.py
@@ -14,20 +14,23 @@
 from py4j.clientserver import ClientServer, JavaParameters, PythonParameters
 from queue import Empty
 from jupyter_client.manager import KernelManager
+from jupyter_client.kernelspec import NoSuchKernel
 import json
 import sys
 
 
-class PythonMagic:
+class JVMKernelMagic:
 
-    def __init__(self):
+    def __init__(self, kernel_name):
         self.km = None
         self.kc = None
         self.comms = []
+        self.kernel_name = kernel_name
+        self.start()
 
     def start(self):
         self.km = KernelManager()
-        self.km.kernel_name = 'python3'
+        self.km.kernel_name = self.kernel_name
         self.km.start_kernel()
         self.kc = self.km.client()
         self.kc.start_channels()
@@ -63,8 +66,8 @@ def pass_msg(self, msg_raw):
 
 class PythonEntryPoint(object):
 
-    def __init__(self):
-        self.pm = PythonMagic()
+    def __init__(self, kernel_name):
+        self.pm = JVMKernelMagic(kernel_name)
 
     def evaluate(self, code):
         print('code for evaluate {}'.format(code))
@@ -92,14 +95,17 @@ class Java:
 
 
 class Py4JServer:
-    def __init__(self, port, pyport):
-        pep = PythonEntryPoint()
-        gateway = ClientServer(
+    def __init__(self, port, pyport, kernel_name):
+        try:
+            pep = PythonEntryPoint(kernel_name)
+        except NoSuchKernel:
+            sys.exit(2)
+        ClientServer(
             java_parameters=JavaParameters(port=int(port)),
             python_parameters=PythonParameters(port=int(pyport)),
             python_server_entry_point=pep)
         print('Py4j server is running')
 
 
 if __name__ == '__main__':
-    Py4JServer(sys.argv[1], sys.argv[2])
+    Py4JServer(sys.argv[1], sys.argv[2], sys.argv[3])
diff --git a/kernel/base/src/main/java/com/twosigma/beakerx/kernel/Kernel.java b/kernel/base/src/main/java/com/twosigma/beakerx/kernel/Kernel.java
@@ -39,6 +39,7 @@
 
 import java.nio.file.Path;
 import java.util.ArrayList;
+import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
 import java.util.Set;
@@ -51,7 +52,7 @@ public abstract class Kernel implements KernelFunctionality {
 
   private static final Logger logger = LoggerFactory.getLogger(Kernel.class);
 
-  public static String OS = System.getProperty("os.name").toLowerCase();
+  private static String OS = System.getProperty("os.name").toLowerCase();
   public static boolean showNullExecutionResult = true;
   private final CloseKernelAction closeKernelAction;
 
@@ -65,7 +66,8 @@ public abstract class Kernel implements KernelFunctionality {
   private List<MagicCommandType> magicCommandTypes;
   private CacheFolderFactory cacheFolderFactory;
   private CustomMagicCommandsFactory customMagicCommands;
-  private PythonMagicManager pythonMagicManager;
+  private Map<String, MagicKernelManager> magicKernels;
+  private Map<String, String> commKernelMapping;
 
   public Kernel(final String sessionId, final Evaluator evaluator,
                 final KernelSocketsFactory kernelSocketsFactory, CustomMagicCommandsFactory customMagicCommands) {
@@ -87,7 +89,8 @@ protected Kernel(final String sessionId, final Evaluator evaluator, final Kernel
     this.executionResultSender = new ExecutionResultSender(this);
     this.evaluatorManager = new EvaluatorManager(this, evaluator);
     this.handlers = new KernelHandlers(this, getCommOpenHandler(this), getKernelInfoHandler(this));
-    this.pythonMagicManager = new PythonMagicManager();
+    this.magicKernels = new HashMap<>();
+    this.commKernelMapping = new HashMap<>();
     createMagicCommands();
     DisplayerDataMapper.init();
     configureSignalHandler();
@@ -115,19 +118,25 @@ public void run() {
   }
 
   private void doExit() {
-    this.pythonMagicManager.exit();
+    doExitOnKernelManager();
     this.evaluatorManager.exit();
     this.handlers.exit();
     this.executionResultSender.exit();
     this.closeKernelAction.close();
   }
 
+  private void doExitOnKernelManager() {
+    for (MagicKernelManager manager : magicKernels.values()) {
+      manager.exit();
+    }
+  }
+
   private void closeComms() {
     this.commMap.values().forEach(Comm::close);
   }
 
   public static boolean isWindows() {
-    return (OS.indexOf("win") >= 0);
+    return OS.contains("win");
   }
 
   public synchronized void setShellOptions(final EvaluatorParameters kernelParameters) {
@@ -292,12 +301,23 @@ public void registerCancelHook(Hook hook) {
   }
 
   @Override
-  public PythonEntryPoint getPythonEntryPoint() {
-      return pythonMagicManager.getPythonEntryPoint();
+  public PythonEntryPoint getPythonEntryPoint(String kernelName) throws NoSuchKernelException {
+      MagicKernelManager manager = magicKernels.get(kernelName);
+      if (manager == null) {
+        manager = new MagicKernelManager(kernelName);
+        magicKernels.put(kernelName, manager);
+      }
+      return manager.getPythonEntryPoint();
+  }
+
+  @Override
+  public MagicKernelManager getManagerByCommId(String commId) {
+    String kernelName = commKernelMapping.get(commId);
+    return magicKernels.get(kernelName);
   }
 
   @Override
-  public PythonMagicManager getPythonMagicManager() {
-    return pythonMagicManager;
+  public void addCommIdManagerMapping(String commId, String kernel){
+    commKernelMapping.put(commId, kernel);
   }
 }
diff --git a/kernel/base/src/main/java/com/twosigma/beakerx/kernel/KernelFunctionality.java b/kernel/base/src/main/java/com/twosigma/beakerx/kernel/KernelFunctionality.java
@@ -25,7 +25,6 @@
 import com.twosigma.beakerx.kernel.magic.command.MagicCommandType;
 import com.twosigma.beakerx.kernel.msg.JupyterMessages;
 import com.twosigma.beakerx.message.Message;
-import py4j.ClientServer;
 
 import java.nio.file.Path;
 import java.util.List;
@@ -94,7 +93,9 @@ public interface KernelFunctionality {
 
   void registerCancelHook(Hook hook);
 
-  PythonEntryPoint getPythonEntryPoint();
+  PythonEntryPoint getPythonEntryPoint(String kernelName) throws NoSuchKernelException;
 
-  PythonMagicManager getPythonMagicManager();
+  MagicKernelManager getManagerByCommId(String commId);
+
+  void addCommIdManagerMapping(String commId, String kernel);
 }
diff --git a/...ma/beakerx/kernel/PythonMagicManager.java → ...ma/beakerx/kernel/MagicKernelManager.java b/...ma/beakerx/kernel/PythonMagicManager.java → ...ma/beakerx/kernel/MagicKernelManager.java
@@ -32,19 +32,27 @@
 import java.util.List;
 import java.util.Map;
 
-public class PythonMagicManager {
+public class MagicKernelManager {
 
     ClientServer clientServer = null;
     private PythonEntryPoint pep = null;
     private Process pythonProcess = null;
 
     private static String DEFAULT_PORT = "25333";
     private static String DEFAULT_PYTHON_PORT = "25334";
+    private static int NO_SUCH_KERNEL_CODE = 2;
+    private static String PY4J_INIT_MESSAGE = "Py4j server is running";
 
     private Integer port = null;
     private Integer pythonPort = null;
 
-    private void initPythonProcess() {
+    private final String kernelName;
+
+    public MagicKernelManager(String kernelName) {
+        this.kernelName = kernelName;
+    }
+
+    private void initPythonProcess() throws NoSuchKernelException {
         //cleanup communication resources if already in use
         exit();
 
@@ -55,8 +63,13 @@ private void initPythonProcess() {
             ProcessBuilder pb = new ProcessBuilder(getPy4jCommand());
             pb.redirectError(ProcessBuilder.Redirect.INHERIT);
             pythonProcess = pb.start();
-            //wait for python process to initialize properly
-            new BufferedReader(new InputStreamReader(pythonProcess.getInputStream())).readLine();
+            BufferedReader br = new BufferedReader(new InputStreamReader(pythonProcess.getInputStream()));
+            while (!PY4J_INIT_MESSAGE.equals(br.readLine()) && pythonProcess.isAlive()){
+                //wait for python process to initialize properly
+            }
+            if (!pythonProcess.isAlive() && pythonProcess.exitValue() == NO_SUCH_KERNEL_CODE) {
+                throw new NoSuchKernelException(kernelName);
+            }
         } catch (IOException e) {
             e.printStackTrace();
         }
@@ -67,7 +80,8 @@ protected String[] getPy4jCommand() {
                 "beakerx",
                 "py4j_server",
                 "--port", port == null ? DEFAULT_PORT : String.valueOf(port),
-                "--pyport", pythonPort == null ? DEFAULT_PYTHON_PORT : String.valueOf(pythonPort)
+                "--pyport", pythonPort == null ? DEFAULT_PYTHON_PORT : String.valueOf(pythonPort),
+                "--kernel", kernelName
         };
     }
 
@@ -86,7 +100,7 @@ public void exit() {
         }
     }
 
-    public PythonEntryPoint getPythonEntryPoint() {
+    public PythonEntryPoint getPythonEntryPoint() throws NoSuchKernelException {
         if (pythonProcess == null || !pythonProcess.isAlive() || clientServer == null) {
             initPythonProcess();
         }
@@ -133,9 +147,13 @@ private Message parseMessage(String stringJson) throws IOException {
     }
 
     public List<Message> handleMsg(Message message) {
-        getPythonEntryPoint();
-        pep.sendMessage(new ObjectMapper().valueToTree(message).toString());
         List<Message> messages = new ArrayList<>();
+        try {
+            getPythonEntryPoint();
+        } catch (NoSuchKernelException e) {
+            return messages;
+        }
+        pep.sendMessage(new ObjectMapper().valueToTree(message).toString());
 
         try {
             messages = getIopubMessages();

diff --git a/kernel/base/src/main/java/com/twosigma/beakerx/kernel/NoSuchKernelException.java b/kernel/base/src/main/java/com/twosigma/beakerx/kernel/NoSuchKernelException.java
@@ -0,0 +1,26 @@
+/*
+ *  Copyright 2017 TWO SIGMA OPEN SOURCE, LLC
+ *
+ *  Licensed under the Apache License, Version 2.0 (the "License");
+ *  you may not use this file except in compliance with the License.
+ *  You may obtain a copy of the License at
+ *
+ *         http://www.apache.org/licenses/LICENSE-2.0
+ *
+ *  Unless required by applicable law or agreed to in writing, software
+ *  distributed under the License is distributed on an "AS IS" BASIS,
+ *  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ *  See the License for the specific language governing permissions and
+ *  limitations under the License.
+ */
+package com.twosigma.beakerx.kernel;
+
+public class NoSuchKernelException extends Exception {
+
+    private static final long serialVersionUID = 1L;
+
+    public NoSuchKernelException(String message) {
+        super(message);
+    }
+
+}
diff --git a/kernel/base/src/main/java/com/twosigma/beakerx/kernel/handler/CommMsgHandler.java b/kernel/base/src/main/java/com/twosigma/beakerx/kernel/handler/CommMsgHandler.java
@@ -17,7 +17,7 @@
 
 import com.twosigma.beakerx.handler.KernelHandler;
 import com.twosigma.beakerx.kernel.KernelFunctionality;
-import com.twosigma.beakerx.kernel.PythonMagicManager;
+import com.twosigma.beakerx.kernel.MagicKernelManager;
 import com.twosigma.beakerx.kernel.comm.Comm;
 import com.twosigma.beakerx.message.Message;
 import org.slf4j.Logger;
@@ -46,14 +46,15 @@ public void handle(Message message) {
 
   private void handleMsg(Message message) {
     Map<String, Serializable> commMap = message.getContent();
-    Comm comm = kernel.getComm(getString(commMap, COMM_ID));
+    String commId = getString(commMap, COMM_ID);
+    Comm comm = kernel.getComm(commId);
     logger.debug("Comm message handling, target name: " + (comm != null ? comm.getTargetName() : "undefined"));
     if (comm != null) {
       comm.handleMsg(message);
     } else {
-      PythonMagicManager pythonMagicManager = kernel.getPythonMagicManager();
-      if (pythonMagicManager != null) {
-        List<Message> messages = pythonMagicManager.handleMsg(message);
+      MagicKernelManager magicKernelManager = kernel.getManagerByCommId(commId);
+      if (magicKernelManager != null) {
+        List<Message> messages = magicKernelManager.handleMsg(message);
         if (!messages.isEmpty()) {
           kernel.publish(messages);
           return;