Create specialized kernels w/ JDK, JVM args, and environment variables (#287)

* Use KOTLIN_JUPYTER_JAVA_HOME or JAVA_HOME if set

* Add possibility to add kernels with different run options

Author: rnett

distrib/kotlin_kernel/__main__.py

@@ -1,7 +1,17 @@
 from kotlin_kernel.install_user import install_user
+from kotlin_kernel.add_kernel import add_kernel
 
 import sys
 
 if __name__ == "__main__":
     if len(sys.argv) == 2 and sys.argv[1] == "fix-kernelspec-location":
         install_user()
+    elif len(sys.argv) >= 2 and sys.argv[1] == "add-kernel":
+        add_kernel()
+    else:
+        if len(sys.argv) < 2:
+            print("Must specify a command", file=sys.stderr)
+        else:
+            print("Unknown command " + sys.argv[1] + ", known commands are fix-kernelspec-location and add-kernel.",
+                  file=sys.stderr)
+        exit(1)

distrib/kotlin_kernel/add_kernel.py

@@ -0,0 +1,109 @@
+import argparse
+import json
+import os.path
+import platform
+import shutil
+import subprocess
+import sys
+
+from kotlin_kernel import env_names
+from kotlin_kernel.install_user import get_user_jupyter_path
+from kotlin_kernel.install_user import install_base_kernel
+
+
+def add_kernel():
+    parser = argparse.ArgumentParser(
+        prog="add-kernel",
+        description="Add a kernel with specified JDK, JVM args, and environment",
+        fromfile_prefix_chars='@')
+    parser.add_argument("--name",
+                        help="The kernel's sub-name.  The kernel will be named \"Kotlin ($name)\".  "
+                             "Will be autodetected if JDK is specified, otherwise required.  "
+                             "Must be file system compatible.")
+    parser.add_argument("--jdk",
+                        help="The home directory of the JDK to use")
+    parser.add_argument("--jvm-arg", action='append', default=[],
+                        help="Add a JVM argument")
+    parser.add_argument("--env", action='append', nargs=2, default=[],
+                        help="Add an environment variable")
+    parser.add_argument("--set-jvm-args", action="store_true", default=False,
+                        help="Set JVM args instead of adding them.")
+    parser.add_argument("--force", action="store_true", default=False,
+                        help="Overwrite an existing kernel with the same name.")
+
+    if len(sys.argv) == 2:
+        parser.print_usage()
+        exit(0)
+
+    args = parser.parse_args(sys.argv[2:])
+
+    jdk = args.jdk
+    if jdk is not None:
+        jdk = os.path.abspath(os.path.expanduser(jdk))
+
+    name = args.name
+    env = {e[0]: e[1] for e in args.env}
+
+    for arg in [env_names.JAVA_HOME, env_names.KERNEL_JAVA_HOME, env_names.JAVA_OPTS,
+                env_names.KERNEL_EXTRA_JAVA_OPTS, env_names.KERNEL_INTERNAL_ADDED_JAVA_OPTS]:
+        if arg in env:
+            print(
+                "Specified environment variable " + arg + ", will be ignored.  "
+                                                          "Use the corresponding arguments instead.", file=sys.stderr)
+            del env[arg]
+
+    if args.set_jvm_args:
+        env[env_names.KERNEL_JAVA_OPTS] = " ".join(args.jvm_arg)
+    else:
+        env[env_names.KERNEL_INTERNAL_ADDED_JAVA_OPTS] = " ".join(args.jvm_arg)
+
+    if jdk is not None:
+        env[env_names.KERNEL_JAVA_HOME] = jdk
+        if platform.system() == 'Windows':
+            java = os.path.join(jdk, "bin/java.exe")
+        else:
+            java = os.path.join(jdk, "bin/java")
+
+        if not os.path.exists(java):
+            print("JDK " + jdk + " has no bin/" + os.path.basename(java), file=sys.stderr)
+            exit(1)
+
+        if name is None:
+            version_spec = subprocess.check_output([java, "--version"], text=True).splitlines()[0].split(" ")
+            dist = version_spec[0]
+            version = version_spec[1]
+            name = "JDK " + dist + " " + version
+
+    if name is None:
+        print("name is required when JDK not specified.", file=sys.stderr)
+        exit(1)
+
+    kernel_name = "kotlin_" + name.replace(" ", "_")
+    kernel_location = os.path.join(get_user_jupyter_path(), "kernels", kernel_name)
+
+    print("Installing kernel to", kernel_location)
+
+    if os.path.exists(kernel_location):
+        if args.force:
+            print("Overwriting existing kernel at " + kernel_location, file=sys.stderr)
+            shutil.rmtree(kernel_location)
+        else:
+            print("There is already a kernel with name " + kernel_name + ", specify a different name "
+                                                                         "or use --force to overwrite it",
+                  file=sys.stderr)
+            exit(1)
+
+    install_base_kernel(kernel_name)
+
+    with open(os.path.join(kernel_location, "kernel.json")) as kernel_file:
+        kernelspec = json.load(kernel_file)
+
+    kernelspec["display_name"] = "Kotlin (" + name + ")"
+
+    if "env" in kernelspec:
+        kernelspec["env"].update(env)
+    else:
+        kernelspec["env"] = env
+
+    with open(os.path.join(kernel_location, "kernel.json"), "w") as kernel_file:
+        json.dump(kernelspec, kernel_file, indent=4)

distrib/kotlin_kernel/env_names.py

@@ -0,0 +1,17 @@
+# standard JVM options environment variable
+JAVA_OPTS = "JAVA_OPTS"
+
+# specific JVM options environment variable
+KERNEL_JAVA_OPTS = "KOTLIN_JUPYTER_JAVA_OPTS"
+
+# additional JVM options to add to either JAVA_OPTS or KOTLIN_JUPYTER_JAVA_OPTS
+KERNEL_EXTRA_JAVA_OPTS = "KOTLIN_JUPYTER_JAVA_OPTS_EXTRA"
+
+# used internally to add JVM options without overwriting KOTLIN_JUPYTER_JAVA_OPTS_EXTRA
+KERNEL_INTERNAL_ADDED_JAVA_OPTS = "KOTLIN_JUPYTER_KERNEL_EXTRA_JVM_OPTS"
+
+# standard JDK location environment variable
+JAVA_HOME = "JAVA_HOME"
+
+# specific JDK location environment variable
+KERNEL_JAVA_HOME = "KOTLIN_JUPYTER_JAVA_HOME"

distrib/kotlin_kernel/install_user.py

@@ -1,29 +1,40 @@
+import os.path
 import platform
 import shutil
 import site
 import sys
 from os import path, environ
 
 
-def install_user():
-    data_relative_path = 'share/jupyter/kernels/kotlin'
-    user_location = path.join(site.getuserbase(), data_relative_path)
-    sys_location = path.join(sys.prefix, data_relative_path)
-    src_paths = [user_location, sys_location]
-
+def get_user_jupyter_path() -> str:
     platform_name = platform.system()
 
     if platform_name == 'Linux':
-        user_jupyter_path = '~/.local/share/jupyter'
+        jupyter_path = '~/.local/share/jupyter'
     elif platform_name == 'Darwin':
-        user_jupyter_path = '~/Library/Jupyter'
+        jupyter_path = '~/Library/Jupyter'
     elif platform_name == 'Windows':
-        user_jupyter_path = path.join(environ['APPDATA'], 'jupyter')
+        jupyter_path = path.join(environ['APPDATA'], 'jupyter')
     else:
         raise OSError("Unknown platform: " + platform_name)
 
-    dst = path.join(user_jupyter_path, 'kernels/kotlin')
+    return os.path.abspath(os.path.expanduser(jupyter_path))
+
+
+def install_base_kernel(kernel_name: str):
+    data_relative_path = 'share/jupyter/kernels/kotlin'
+    user_location = path.join(site.getuserbase(), data_relative_path)
+    sys_location = path.join(sys.prefix, data_relative_path)
+    src_paths = [user_location, sys_location]
+
+    user_jupyter_path = get_user_jupyter_path()
+
+    dst = path.join(user_jupyter_path, 'kernels/' + kernel_name)
     for src in src_paths:
         if not path.exists(src):
             continue
         shutil.copytree(src, dst, dirs_exist_ok=True)
+
+
+def install_user():
+    install_base_kernel('kotlin')

distrib/run_kotlin_kernel/run_kernel.py

@@ -1,9 +1,12 @@
 import json
 import os
+import shlex
 import subprocess
 import sys
 from typing import List
 
+from kotlin_kernel import env_names
+
 
 def run_kernel(*args) -> None:
     try:
@@ -39,7 +42,25 @@ def run_kernel_impl(connection_file: str, jar_args_file: str = None, executables
         class_path_arg = os.pathsep.join([os.path.join(jars_dir, jar_name) for jar_name in cp])
         main_jar_path = os.path.join(jars_dir, main_jar)
 
-        subprocess.call(['java', '-jar'] + debug_list +
+        java_home = os.getenv(env_names.KERNEL_JAVA_HOME) or os.getenv(env_names.JAVA_HOME)
+
+        if java_home is None:
+            java = "java"
+        else:
+            java = os.path.join(java_home, "bin", "java")
+
+        jvm_arg_str = os.getenv(env_names.KERNEL_JAVA_OPTS) or os.getenv(env_names.JAVA_OPTS) or ""
+        extra_args = os.getenv(env_names.KERNEL_EXTRA_JAVA_OPTS)
+        if extra_args is not None:
+            jvm_arg_str += " " + extra_args
+
+        kernel_args = os.getenv(env_names.KERNEL_INTERNAL_ADDED_JAVA_OPTS)
+        if kernel_args is not None:
+            jvm_arg_str += " " + kernel_args
+
+        jvm_args = shlex.split(jvm_arg_str)
+
+        subprocess.call([java] + jvm_args + ['-jar'] + debug_list +
                         [main_jar_path,
                          '-classpath=' + class_path_arg,
                          connection_file,

docs/README-STUB.md

@@ -84,6 +84,39 @@ Don't forget to re-run this script on the kernel update.
 
 To start using `kotlin` kernel inside Jupyter Notebook or JupyterLab create a new notebook with `kotlin` kernel.
 
+The default kernel will use the JDK pointed to by the environment variable `KOTLIN_JUPYTER_JAVA_HOME`,
+or `JAVA_HOME` if the first is not set.
+
+JVM arguments will be set from the environment variable `KOTLIN_JUPYTER_JAVA_OPTS` or `JAVA_OPTS` if the first is not set.
+Additionally, arguments from `KOTLIN_JUPYTER_JAVA_OPTS_EXTRA` will be added.
+Arguments are parsed using [`shlex.split`](https://docs.python.org/3/library/shlex.html).
+
+### Creating Kernels
+
+To create a kernel for a specific JDK, JVM arguments, and environment variables, you can use the `add-kernel` script:
+```bash
+python -m kotlin_kernel add-kernel [--name name] [--jdk jdk_home_dir] [--set-jvm-args] [--jvm-arg arg]* [--env KEY VALUE]* [--force]
+```
+The command uses `argparse`, so `--help`, `@argfile` (you will need to escape the `@` in powershell), and `--opt=value` are all supported.  `--jvm-arg=arg` in particular
+is needed when passing JVM arguments that start with `-`.
+
+If `jdk` not specified, `name` is required.  If `name` is not specified but `jdk` is the name will be 
+`JDK $vendor $version` detected from the JDK.  Regardless, the actual name of the kernel will be `Kotlin ($name)`, 
+and the directory will be `kotlin_$name` with the spaces in `name` replaced by underscores 
+(so make sure it's compatible with your file system).
+
+JVM arguments are joined with a `' '`, so multiple JVM arguments in the same argument are supported.
+The arguments will be added to existing ones (see above section) unless `--set-jvm-args` is present, in which case they
+will be set to `KOTLIN_JUPYTER_JAVA_OPTS`.  Note that both adding and setting work fine alongside `KOTLIN_JUPYTER_JAVA_OPTS_EXTRA`.
+
+While jupyter kernel environment variable substitutions are supported in `env`, note that if the used environment 
+variable doesn't exist, nothing will be replaced.
+
+An example:
+```bash
+python -m kotlin_kernel add-kernel --name "JDK 15 Big 2 GPU" --jdk ~/.jdks/openjdk-15.0.2 --jvm-arg=-Xmx8G --env CUDA_VISIBLE_DEVICES 0,1
+```
+
 ## Supported functionality
 
 ### REPL commands

docs/README.md

@@ -84,6 +84,39 @@ Don't forget to re-run this script on the kernel update.
 
 To start using `kotlin` kernel inside Jupyter Notebook or JupyterLab create a new notebook with `kotlin` kernel.
 
+The default kernel will use the JDK pointed to by the environment variable `KOTLIN_JUPYTER_JAVA_HOME`,
+or `JAVA_HOME` if the first is not set.
+
+JVM arguments will be set from the environment variable `KOTLIN_JUPYTER_JAVA_OPTS` or `JAVA_OPTS` if the first is not set.
+Additionally, arguments from `KOTLIN_JUPYTER_JAVA_OPTS_EXTRA` will be added.
+Arguments are parsed using [`shlex.split`](https://docs.python.org/3/library/shlex.html).
+
+### Creating Kernels
+
+To create a kernel for a specific JDK, JVM arguments, and environment variables, you can use the `add-kernel` script:
+```bash
+python -m kotlin_kernel add-kernel [--name name] [--jdk jdk_home_dir] [--set-jvm-args] [--jvm-arg arg]* [--env KEY VALUE]* [--force]
+```
+The command uses `argparse`, so `--help`, `@argfile` (you will need to escape the `@` in powershell), and `--opt=value` are all supported.  `--jvm-arg=arg` in particular
+is needed when passing JVM arguments that start with `-`.
+
+If `jdk` not specified, `name` is required.  If `name` is not specified but `jdk` is the name will be 
+`JDK $vendor $version` detected from the JDK.  Regardless, the actual name of the kernel will be `Kotlin ($name)`, 
+and the directory will be `kotlin_$name` with the spaces in `name` replaced by underscores 
+(so make sure it's compatible with your file system).
+
+JVM arguments are joined with a `' '`, so multiple JVM arguments in the same argument are supported.
+The arguments will be added to existing ones (see above section) unless `--set-jvm-args` is present, in which case they
+will be set to `KOTLIN_JUPYTER_JAVA_OPTS`.  Note that both adding and setting work fine alongside `KOTLIN_JUPYTER_JAVA_OPTS_EXTRA`.
+
+While jupyter kernel environment variable substitutions are supported in `env`, note that if the used environment 
+variable doesn't exist, nothing will be replaced.
+
+An example:
+```bash
+python -m kotlin_kernel add-kernel --name "JDK 15 Big 2 GPU" --jdk ~/.jdks/openjdk-15.0.2 --jvm-arg=-Xmx8G --env CUDA_VISIBLE_DEVICES 0,1
+```
+
 ## Supported functionality
 
 ### REPL commands