Feature: JupyterLab plugin

aztk/models/plugins/internal/plugin_manager.py

@@ -17,6 +17,7 @@ class PluginManager:
     # Indexing of all the predefined plugins
     plugins = dict(
         jupyter=plugins.JupyterPlugin,
+        jupyter_lab=plugins.JupyterLabPlugin,
         rstudio_server=plugins.RStudioServerPlugin,
         hdfs=plugins.HDFSPlugin,
     )

aztk/spark/models/plugins/__init__.py

@@ -1,3 +1,4 @@
 from .hdfs import *
 from .jupyter import *
+from .jupyter_lab import *
 from .rstudio_server import *

aztk/spark/models/plugins/jupyter_lab/__init__.py

@@ -0,0 +1 @@
+from .configuration import *

aztk/spark/models/plugins/jupyter_lab/configuration.py

@@ -0,0 +1,23 @@
+import os
+from aztk.models.plugins.plugin_configuration import PluginConfiguration, PluginPort, PluginRunTarget
+from aztk.models.plugins.plugin_file import PluginFile
+from aztk.utils import constants
+
+dir_path = os.path.dirname(os.path.realpath(__file__))
+
+class JupyterLabPlugin(PluginConfiguration):
+    def __init__(self):
+        super().__init__(
+            name="jupyter_lab",
+            ports=[
+                PluginPort(
+                    internal=8889,
+                    public=True,
+                ),
+            ],
+            run_on=PluginRunTarget.All,
+            execute="jupyter_lab.sh",
+            files=[
+                PluginFile("jupyter_lab.sh", os.path.join(dir_path, "jupyter_lab.sh")),
+            ],
+        )

aztk/spark/models/plugins/jupyter_lab/jupyter_lab.sh

@@ -0,0 +1,57 @@
+#!/bin/bash
+
+# This custom script only works on images where jupyter is pre-installed on the Docker image
+#
+# This custom script has been tested to work on the following docker images:
+#  - aztk/python:spark2.2.0-python3.6.2-base
+#  - aztk/python:spark2.2.0-python3.6.2-gpu
+#  - aztk/python:spark2.1.0-python3.6.2-base
+#  - aztk/python:spark2.1.0-python3.6.2-gpu
+
+if  [ "$IS_MASTER" = "1" ]; then
+    conda install -c conda-force jupyterlab
+
+    PYSPARK_DRIVER_PYTHON="/.pyenv/versions/${USER_PYTHON_VERSION}/bin/jupyter"
+    JUPYTER_KERNELS="/.pyenv/versions/${USER_PYTHON_VERSION}/share/jupyter/kernels"
+
+    # disable password/token on jupyter notebook
+    jupyter lab --generate-config --allow-root
+    JUPYTER_CONFIG='/.jupyter/jupyter_notebook_config.py'
+    echo >> $JUPYTER_CONFIG
+    echo -e 'c.NotebookApp.token=""' >> $JUPYTER_CONFIG
+    echo -e 'c.NotebookApp.password=""' >> $JUPYTER_CONFIG
+
+    # get master ip
+    MASTER_IP=$(hostname -i)
+
+    # remove existing kernels
+    rm -rf $JUPYTER_KERNELS/*
+
+    # set up jupyter to use pyspark
+    mkdir $JUPYTER_KERNELS/pyspark
+    touch $JUPYTER_KERNELS/pyspark/kernel.json
+    cat << EOF > $JUPYTER_KERNELS/pyspark/kernel.json
+{
+    "display_name": "PySpark",
+    "language": "python",
+    "argv": [
+        "python",
+        "-m",
+        "ipykernel",
+        "-f",
+        "{connection_file}"
+    ],
+    "env": {
+        "SPARK_HOME": "$SPARK_HOME",
+        "PYSPARK_PYTHON": "python",
+        "PYSPARK_SUBMIT_ARGS": "--master spark://$MASTER_IP:7077 pyspark-shell"
+    }
+}
+EOF
+
+    # start jupyter notebook from /mnt - this is where we recommend you put your azure files mount point as well
+    cd /mnt
+    (PYSPARK_DRIVER_PYTHON=$PYSPARK_DRIVER_PYTHON PYSPARK_DRIVER_PYTHON_OPTS="lab --no-browser --port=8889 --allow-root" pyspark &)
+fi
+
+

node_scripts/install/install.py

@@ -27,6 +27,11 @@ def setup_node():
         plugins.setup_plugins(is_master=True, is_worker=True)
         scripts.run_custom_scripts(is_master=True, is_worker=True)
 
+        #Write sentinel file on disk for host to know this is the master
+        path = os.environ['PWD'] + '/MASTER'
+        print('Writing master file to: {}'.format(path))
+        f = open(path, 'w')
+        f.write(master_node.ip_address)
     else:
         setup_as_worker()
         plugins.setup_plugins(is_master=False, is_worker=True)
@@ -42,7 +47,6 @@ def setup_as_master():
     if os.environ["WORKER_ON_MASTER"] == "True":
         spark.start_spark_worker()
 
-
 def setup_as_worker():
     print("Setting up as worker.")
     spark.setup_connection()

node_scripts/setup_node.sh

@@ -4,23 +4,26 @@
 # Usage:
 # setup_node.sh [container_name] [gpu_enabled] [docker_repo] [docker_cmd]
 
-
 container_name=$1
 gpu_enabled=$2
 repo_name=$3
 docker_run_cmd=$4
 
+echo "Installing pre-reqs"
 apt-get -y install linux-image-extra-$(uname -r) linux-image-extra-virtual
 apt-get -y install apt-transport-https
 apt-get -y install curl
 apt-get -y install ca-certificates
 apt-get -y install software-properties-common
+echo "Done installing pre-reqs"
 
 # Install docker
+echo "Installing Docker"
 curl -fsSL https://download.docker.com/linux/ubuntu/gpg | sudo apt-key add -
 add-apt-repository "deb [arch=amd64] https://download.docker.com/linux/ubuntu $(lsb_release -cs) stable"
 apt-get -y update
 apt-get -y install docker-ce
+echo "Done installing Docker"
 
 if [ $gpu_enabled == "True" ]; then
     echo "running nvidia install"
@@ -71,7 +74,7 @@ else
     until [ "`/usr/bin/docker inspect -f {{.State.Running}} $container_name`"=="true" ]; do
         sleep 0.1;
     done;
-    
+
     # wait until container setup is complete
     docker exec spark /bin/bash -c 'python $DOCKER_WORKING_DIR/wait_until_setup_complete.py'