Skip to content

Commit 36ff53e

Browse files
rnettrnett
authored
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
1 parent fd50133 commit 36ff53e

File tree

7 files changed

+245
-11
lines changed

7 files changed

+245
-11
lines changed

‎distrib/kotlin_kernel/__main__.py

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,17 @@
11
from kotlin_kernel.install_user import install_user
2+
from kotlin_kernel.add_kernel import add_kernel
23

34
import sys
45

56
if __name__ == "__main__":
67
if len(sys.argv) == 2 and sys.argv[1] == "fix-kernelspec-location":
78
install_user()
9+
elif len(sys.argv) >= 2 and sys.argv[1] == "add-kernel":
10+
add_kernel()
11+
else:
12+
if len(sys.argv) < 2:
13+
print("Must specify a command", file=sys.stderr)
14+
else:
15+
print("Unknown command " + sys.argv[1] + ", known commands are fix-kernelspec-location and add-kernel.",
16+
file=sys.stderr)
17+
exit(1)

‎distrib/kotlin_kernel/add_kernel.py

Lines changed: 109 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,109 @@
1+
import argparse
2+
import json
3+
import os.path
4+
import platform
5+
import shutil
6+
import subprocess
7+
import sys
8+
9+
from kotlin_kernel import env_names
10+
from kotlin_kernel.install_user import get_user_jupyter_path
11+
from kotlin_kernel.install_user import install_base_kernel
12+
13+
14+
def add_kernel():
15+
parser = argparse.ArgumentParser(
16+
prog="add-kernel",
17+
description="Add a kernel with specified JDK, JVM args, and environment",
18+
fromfile_prefix_chars='@')
19+
parser.add_argument("--name",
20+
help="The kernel's sub-name. The kernel will be named \"Kotlin ($name)\". "
21+
"Will be autodetected if JDK is specified, otherwise required. "
22+
"Must be file system compatible.")
23+
parser.add_argument("--jdk",
24+
help="The home directory of the JDK to use")
25+
parser.add_argument("--jvm-arg", action='append', default=[],
26+
help="Add a JVM argument")
27+
parser.add_argument("--env", action='append', nargs=2, default=[],
28+
help="Add an environment variable")
29+
parser.add_argument("--set-jvm-args", action="store_true", default=False,
30+
help="Set JVM args instead of adding them.")
31+
parser.add_argument("--force", action="store_true", default=False,
32+
help="Overwrite an existing kernel with the same name.")
33+
34+
if len(sys.argv) == 2:
35+
parser.print_usage()
36+
exit(0)
37+
38+
args = parser.parse_args(sys.argv[2:])
39+
40+
jdk = args.jdk
41+
if jdk is not None:
42+
jdk = os.path.abspath(os.path.expanduser(jdk))
43+
44+
name = args.name
45+
env = {e[0]: e[1] for e in args.env}
46+
47+
for arg in [env_names.JAVA_HOME, env_names.KERNEL_JAVA_HOME, env_names.JAVA_OPTS,
48+
env_names.KERNEL_EXTRA_JAVA_OPTS, env_names.KERNEL_INTERNAL_ADDED_JAVA_OPTS]:
49+
if arg in env:
50+
print(
51+
"Specified environment variable " + arg + ", will be ignored. "
52+
"Use the corresponding arguments instead.", file=sys.stderr)
53+
del env[arg]
54+
55+
if args.set_jvm_args:
56+
env[env_names.KERNEL_JAVA_OPTS] = " ".join(args.jvm_arg)
57+
else:
58+
env[env_names.KERNEL_INTERNAL_ADDED_JAVA_OPTS] = " ".join(args.jvm_arg)
59+
60+
if jdk is not None:
61+
env[env_names.KERNEL_JAVA_HOME] = jdk
62+
if platform.system() == 'Windows':
63+
java = os.path.join(jdk, "bin/java.exe")
64+
else:
65+
java = os.path.join(jdk, "bin/java")
66+
67+
if not os.path.exists(java):
68+
print("JDK " + jdk + " has no bin/" + os.path.basename(java), file=sys.stderr)
69+
exit(1)
70+
71+
if name is None:
72+
version_spec = subprocess.check_output([java, "--version"], text=True).splitlines()[0].split(" ")
73+
dist = version_spec[0]
74+
version = version_spec[1]
75+
name = "JDK " + dist + " " + version
76+
77+
if name is None:
78+
print("name is required when JDK not specified.", file=sys.stderr)
79+
exit(1)
80+
81+
kernel_name = "kotlin_" + name.replace(" ", "_")
82+
kernel_location = os.path.join(get_user_jupyter_path(), "kernels", kernel_name)
83+
84+
print("Installing kernel to", kernel_location)
85+
86+
if os.path.exists(kernel_location):
87+
if args.force:
88+
print("Overwriting existing kernel at " + kernel_location, file=sys.stderr)
89+
shutil.rmtree(kernel_location)
90+
else:
91+
print("There is already a kernel with name " + kernel_name + ", specify a different name "
92+
"or use --force to overwrite it",
93+
file=sys.stderr)
94+
exit(1)
95+
96+
install_base_kernel(kernel_name)
97+
98+
with open(os.path.join(kernel_location, "kernel.json")) as kernel_file:
99+
kernelspec = json.load(kernel_file)
100+
101+
kernelspec["display_name"] = "Kotlin (" + name + ")"
102+
103+
if "env" in kernelspec:
104+
kernelspec["env"].update(env)
105+
else:
106+
kernelspec["env"] = env
107+
108+
with open(os.path.join(kernel_location, "kernel.json"), "w") as kernel_file:
109+
json.dump(kernelspec, kernel_file, indent=4)

‎distrib/kotlin_kernel/env_names.py

Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,17 @@
1+
# standard JVM options environment variable
2+
JAVA_OPTS = "JAVA_OPTS"
3+
4+
# specific JVM options environment variable
5+
KERNEL_JAVA_OPTS = "KOTLIN_JUPYTER_JAVA_OPTS"
6+
7+
# additional JVM options to add to either JAVA_OPTS or KOTLIN_JUPYTER_JAVA_OPTS
8+
KERNEL_EXTRA_JAVA_OPTS = "KOTLIN_JUPYTER_JAVA_OPTS_EXTRA"
9+
10+
# used internally to add JVM options without overwriting KOTLIN_JUPYTER_JAVA_OPTS_EXTRA
11+
KERNEL_INTERNAL_ADDED_JAVA_OPTS = "KOTLIN_JUPYTER_KERNEL_EXTRA_JVM_OPTS"
12+
13+
# standard JDK location environment variable
14+
JAVA_HOME = "JAVA_HOME"
15+
16+
# specific JDK location environment variable
17+
KERNEL_JAVA_HOME = "KOTLIN_JUPYTER_JAVA_HOME"

‎distrib/kotlin_kernel/install_user.py

Lines changed: 21 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -1,29 +1,40 @@
1+
import os.path
12
import platform
23
import shutil
34
import site
45
import sys
56
from os import path, environ
67

78

8-
def install_user():
9-
data_relative_path = 'share/jupyter/kernels/kotlin'
10-
user_location = path.join(site.getuserbase(), data_relative_path)
11-
sys_location = path.join(sys.prefix, data_relative_path)
12-
src_paths = [user_location, sys_location]
13-
9+
def get_user_jupyter_path() -> str:
1410
platform_name = platform.system()
1511

1612
if platform_name == 'Linux':
17-
user_jupyter_path = '~/.local/share/jupyter'
13+
jupyter_path = '~/.local/share/jupyter'
1814
elif platform_name == 'Darwin':
19-
user_jupyter_path = '~/Library/Jupyter'
15+
jupyter_path = '~/Library/Jupyter'
2016
elif platform_name == 'Windows':
21-
user_jupyter_path = path.join(environ['APPDATA'], 'jupyter')
17+
jupyter_path = path.join(environ['APPDATA'], 'jupyter')
2218
else:
2319
raise OSError("Unknown platform: " + platform_name)
2420

25-
dst = path.join(user_jupyter_path, 'kernels/kotlin')
21+
return os.path.abspath(os.path.expanduser(jupyter_path))
22+
23+
24+
def install_base_kernel(kernel_name: str):
25+
data_relative_path = 'share/jupyter/kernels/kotlin'
26+
user_location = path.join(site.getuserbase(), data_relative_path)
27+
sys_location = path.join(sys.prefix, data_relative_path)
28+
src_paths = [user_location, sys_location]
29+
30+
user_jupyter_path = get_user_jupyter_path()
31+
32+
dst = path.join(user_jupyter_path, 'kernels/' + kernel_name)
2633
for src in src_paths:
2734
if not path.exists(src):
2835
continue
2936
shutil.copytree(src, dst, dirs_exist_ok=True)
37+
38+
39+
def install_user():
40+
install_base_kernel('kotlin')

‎distrib/run_kotlin_kernel/run_kernel.py

Lines changed: 22 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,9 +1,12 @@
11
import json
22
import os
3+
import shlex
34
import subprocess
45
import sys
56
from typing import List
67

8+
from kotlin_kernel import env_names
9+
710

811
def run_kernel(*args) -> None:
912
try:
@@ -39,7 +42,25 @@ def run_kernel_impl(connection_file: str, jar_args_file: str = None, executables
3942
class_path_arg = os.pathsep.join([os.path.join(jars_dir, jar_name) for jar_name in cp])
4043
main_jar_path = os.path.join(jars_dir, main_jar)
4144

42-
subprocess.call(['java', '-jar'] + debug_list +
45+
java_home = os.getenv(env_names.KERNEL_JAVA_HOME) or os.getenv(env_names.JAVA_HOME)
46+
47+
if java_home is None:
48+
java = "java"
49+
else:
50+
java = os.path.join(java_home, "bin", "java")
51+
52+
jvm_arg_str = os.getenv(env_names.KERNEL_JAVA_OPTS) or os.getenv(env_names.JAVA_OPTS) or ""
53+
extra_args = os.getenv(env_names.KERNEL_EXTRA_JAVA_OPTS)
54+
if extra_args is not None:
55+
jvm_arg_str += " " + extra_args
56+
57+
kernel_args = os.getenv(env_names.KERNEL_INTERNAL_ADDED_JAVA_OPTS)
58+
if kernel_args is not None:
59+
jvm_arg_str += " " + kernel_args
60+
61+
jvm_args = shlex.split(jvm_arg_str)
62+
63+
subprocess.call([java] + jvm_args + ['-jar'] + debug_list +
4364
[main_jar_path,
4465
'-classpath=' + class_path_arg,
4566
connection_file,

‎docs/README-STUB.md

Lines changed: 33 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -84,6 +84,39 @@ Don't forget to re-run this script on the kernel update.
8484

8585
To start using `kotlin` kernel inside Jupyter Notebook or JupyterLab create a new notebook with `kotlin` kernel.
8686

87+
The default kernel will use the JDK pointed to by the environment variable `KOTLIN_JUPYTER_JAVA_HOME`,
88+
or `JAVA_HOME` if the first is not set.
89+
90+
JVM arguments will be set from the environment variable `KOTLIN_JUPYTER_JAVA_OPTS` or `JAVA_OPTS` if the first is not set.
91+
Additionally, arguments from `KOTLIN_JUPYTER_JAVA_OPTS_EXTRA` will be added.
92+
Arguments are parsed using [`shlex.split`](https://docs.python.org/3/library/shlex.html).
93+
94+
### Creating Kernels
95+
96+
To create a kernel for a specific JDK, JVM arguments, and environment variables, you can use the `add-kernel` script:
97+
```bash
98+
python -m kotlin_kernel add-kernel [--name name] [--jdk jdk_home_dir] [--set-jvm-args] [--jvm-arg arg]* [--env KEY VALUE]* [--force]
99+
```
100+
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
101+
is needed when passing JVM arguments that start with `-`.
102+
103+
If `jdk` not specified, `name` is required. If `name` is not specified but `jdk` is the name will be
104+
`JDK $vendor $version` detected from the JDK. Regardless, the actual name of the kernel will be `Kotlin ($name)`,
105+
and the directory will be `kotlin_$name` with the spaces in `name` replaced by underscores
106+
(so make sure it's compatible with your file system).
107+
108+
JVM arguments are joined with a `' '`, so multiple JVM arguments in the same argument are supported.
109+
The arguments will be added to existing ones (see above section) unless `--set-jvm-args` is present, in which case they
110+
will be set to `KOTLIN_JUPYTER_JAVA_OPTS`. Note that both adding and setting work fine alongside `KOTLIN_JUPYTER_JAVA_OPTS_EXTRA`.
111+
112+
While jupyter kernel environment variable substitutions are supported in `env`, note that if the used environment
113+
variable doesn't exist, nothing will be replaced.
114+
115+
An example:
116+
```bash
117+
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
118+
```
119+
87120
## Supported functionality
88121

89122
### REPL commands

‎docs/README.md

Lines changed: 33 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -84,6 +84,39 @@ Don't forget to re-run this script on the kernel update.
8484

8585
To start using `kotlin` kernel inside Jupyter Notebook or JupyterLab create a new notebook with `kotlin` kernel.
8686

87+
The default kernel will use the JDK pointed to by the environment variable `KOTLIN_JUPYTER_JAVA_HOME`,
88+
or `JAVA_HOME` if the first is not set.
89+
90+
JVM arguments will be set from the environment variable `KOTLIN_JUPYTER_JAVA_OPTS` or `JAVA_OPTS` if the first is not set.
91+
Additionally, arguments from `KOTLIN_JUPYTER_JAVA_OPTS_EXTRA` will be added.
92+
Arguments are parsed using [`shlex.split`](https://docs.python.org/3/library/shlex.html).
93+
94+
### Creating Kernels
95+
96+
To create a kernel for a specific JDK, JVM arguments, and environment variables, you can use the `add-kernel` script:
97+
```bash
98+
python -m kotlin_kernel add-kernel [--name name] [--jdk jdk_home_dir] [--set-jvm-args] [--jvm-arg arg]* [--env KEY VALUE]* [--force]
99+
```
100+
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
101+
is needed when passing JVM arguments that start with `-`.
102+
103+
If `jdk` not specified, `name` is required. If `name` is not specified but `jdk` is the name will be
104+
`JDK $vendor $version` detected from the JDK. Regardless, the actual name of the kernel will be `Kotlin ($name)`,
105+
and the directory will be `kotlin_$name` with the spaces in `name` replaced by underscores
106+
(so make sure it's compatible with your file system).
107+
108+
JVM arguments are joined with a `' '`, so multiple JVM arguments in the same argument are supported.
109+
The arguments will be added to existing ones (see above section) unless `--set-jvm-args` is present, in which case they
110+
will be set to `KOTLIN_JUPYTER_JAVA_OPTS`. Note that both adding and setting work fine alongside `KOTLIN_JUPYTER_JAVA_OPTS_EXTRA`.
111+
112+
While jupyter kernel environment variable substitutions are supported in `env`, note that if the used environment
113+
variable doesn't exist, nothing will be replaced.
114+
115+
An example:
116+
```bash
117+
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
118+
```
119+
87120
## Supported functionality
88121

89122
### REPL commands

0 commit comments

Comments
 (0)