feature: Add RTC extension and scripts (#277)

Co-authored-by: Zuoyuan Huang <[email protected]>
aws · Apr 4, 2024 · 82c2dec · 82c2dec
1 parent fe0415e
commit 82c2dec
Show file tree

Hide file tree

Showing 6 changed files with 19 additions and 0 deletions.
diff --git a/build_artifacts/v1/v1.7/v1.7.0/cpu.additional_packages_env.in b/build_artifacts/v1/v1.7/v1.7.0/cpu.additional_packages_env.in
@@ -0,0 +1 @@
+conda-forge::jupyter-collaboration[version='>=1.1.0']
diff --git a/build_artifacts/v1/v1.7/v1.7.0/gpu.additional_packages_env.in b/build_artifacts/v1/v1.7/v1.7.0/gpu.additional_packages_env.in
@@ -0,0 +1 @@
+conda-forge::jupyter-collaboration[version='>=1.1.0']
diff --git a/template/v1/Dockerfile b/template/v1/Dockerfile
@@ -6,6 +6,7 @@ ARG ENV_IN_FILENAME
 ARG ARG_BASED_ENV_IN_FILENAME
 
 ARG AMZN_BASE="/opt/amazon/sagemaker"
+ARG DB_ROOT_DIR="/opt/db"
 ARG DIRECTORY_TREE_STAGE_DIR="${AMZN_BASE}/dir-staging"
 
 ARG NB_USER="sagemaker-user"
@@ -108,6 +109,9 @@ RUN install-glue-kernels && \
 COPY patch_glue_pyspark.json /opt/conda/share/jupyter/kernels/glue_pyspark/kernel.json
 COPY patch_glue_spark.json /opt/conda/share/jupyter/kernels/glue_spark/kernel.json
 
+# Configure RTC - disable jupyter_collaboration by default
+RUN jupyter labextension disable @jupyter/collaboration-extension
+
 USER root
 RUN HOME_DIR="/home/${NB_USER}/licenses" \
     && mkdir -p ${HOME_DIR} \
@@ -132,6 +136,11 @@ RUN rm -rf /etc/code-editor
 RUN mkdir -p /var/run/supervisord && \
     chmod a+rw /var/run/supervisord
 
+# Create root directory for DB
+# Create logging directories for supervisor
+RUN mkdir -p $DB_ROOT_DIR && \
+    chmod a+rw $DB_ROOT_DIR
+
 USER $MAMBA_USER
 ENV PATH="/opt/conda/bin:/opt/conda/condabin:$PATH"
 WORKDIR "/home/${NB_USER}"

diff --git a/template/v1/dirs/usr/local/bin/start-jupyter-server b/template/v1/dirs/usr/local/bin/start-jupyter-server
@@ -8,6 +8,7 @@ micromamba activate base
 
 # Start Jupyter server in rtc mode for shared spaces
 if [ -n "$SAGEMAKER_APP_TYPE_LOWERCASE" ] && [ "$SAGEMAKER_SPACE_TYPE_LOWERCASE" == "shared" ]; then
+  jupyter labextension enable @jupyter/collaboration-extension
   # SAGEMAKER_APP_TYPE is set, indicating the server is running within a SageMaker
   # app. Configure the base url to be `/<app-type-in-lower-case>/default`.
   # SAGEMAKER_SPACE_TYPE_LOWERCASE flag is used to determine if the server should start

diff --git a/test/test_artifacts/v1/jupyter-collaboration.test.Dockerfile b/test/test_artifacts/v1/jupyter-collaboration.test.Dockerfile
@@ -0,0 +1,6 @@
+ARG SAGEMAKER_DISTRIBUTION_IMAGE
+FROM $SAGEMAKER_DISTRIBUTION_IMAGE
+
+ARG MAMBA_DOCKERFILE_ACTIVATE=1
+
+CMD ["python", "-c", "import jupyter_collaboration; import jupyter_server_fileid; from jupyter_ydoc import YBlob; yblob = YBlob(); assert yblob.get() == b''; yblob.set(b'012'); assert yblob.get() == b'012'"]
diff --git a/test/test_dockerfile_based_harness.py b/test/test_dockerfile_based_harness.py
@@ -29,6 +29,7 @@
         ("pytorch.examples.Dockerfile", ["pytorch"]),
         ("tensorflow.examples.Dockerfile", ["tensorflow"]),
         ("jupyter-ai.test.Dockerfile", ["jupyter-ai"]),
+        ("jupyter-collaboration.test.Dockerfile", ["jupyter-collaboration"]),
         ("jupyter-dash.test.Dockerfile", ["jupyter-dash"]),
         ("jupyterlab-lsp.test.Dockerfile", ["jupyterlab-lsp"]),
         ("jupyter-lsp-server.test.Dockerfile", ["jupyter-lsp-server"]),
Original file line number	Diff line number	Diff line change
		@@ -0,0 +1 @@
		conda-forge::jupyter-collaboration[version='>=1.1.0']